Economical Jet Taggers -- Equivariant, Slim, and Quantized

Este artigo apresenta uma versão esguia, quantizada e com redução de parâmetros do classificador de jatos L-GATr que alcança uma redução de uma ordem de magnitude no custo energético com apenas uma diminuição moderada de desempenho, pavimentando o caminho para a classificação de jatos eficiente em nível de gatilho no LHC.

O essencial

Imagine o Grande Colisor de Hádrons (LHC) como uma enorme e veloz fábrica de partículas. A cada segundo, ele colide prótons, criando uma dispersão caótica de detritos. Os físicos precisam filtrar esses detritos para encontrar partículas específicas e raras (como o "quark top") escondidas entre bilhões de outras comuns. Esse processo de classificação é chamado de jet tagging.

Por anos, cientistas usaram programas de computador complexos (Aprendizado de Máquina) para realizar essa classificação. Os atuais campeões são os "Transformers" — modelos de IA poderosos que são incrivelmente precisos, mas também enormes, lentos e famintos por energia. Eles são como uma frota de caminhões enormes e gastadores de combustível tentando entregar uma única carta; eles fazem o trabalho, mas são grandes demais e caros demais para serem usados no exato momento em que os dados estão sendo coletados (o nível do "trigger").

Este artigo faz uma pergunta simples: Podemos encolher esses caminhões gigantes em pequenas lambretas econômicas e eficientes sem perder a capacidade de entregar a carta?

Veja como os autores fizeram isso, utilizando três estratégias principais:

1. A Versão "Enxuta" (L-GATr-slim)

O modelo original "L-GATr" é como um canivete suíço que carrega todas as ferramentas possíveis: escalares, vetores, tensores e muito mais. No entanto, os autores perceberam que, para a maioria dos trabalhos de física de partículas, você só precisa de duas ferramentas: escalares (números) e vetores (setas com direção).

• A Analogia: Imagine um chef que insiste em usar uma cozinha industrial completa, com fornos, liquidificadores e batedeiras, apenas para fazer um sanduíche simples. Os autores disseram: "Vamos usar apenas uma faca e uma tábua de corte".
• O Resultado: Eles construíram uma versão "Slim" (enxuta) da IA que remove as ferramentas desnecessárias. Ela tem o mesmo desempenho da versão gigante, mas é muito mais rápida de treinar e usa menos memória. É como trocar um caminhão pesado por um carro esportivo ágil que realiza o mesmo trabalho.

2. A Versão "Minúscula" (Ultra-mini Taggers)

Os autores então perguntaram: "Quão pequeno podemos chegar?" Eles tentaram encolher esses modelos de IA ao tamanho de um carrinho de brinquedo (cerca de 1.000 parâmetros, comparado aos milhões do original).

• A Analogia: Pense em tentar encaixar toda a biblioteca de conhecimento em um único cartão-postal. Geralmente, você perde a história. Mas os autores descobriram que, se organizarem a informação corretamente (usando regras específicas "Lorentz-equivariantes" que respeitam as leis da física), podem encaixar o conhecimento essencial em um espaço minúsculo.
• O Resultado: Eles descobriram que, para modelos muito pequenos, a arquitetura "LLoCa" funciona melhor se você encolher o número de camadas, enquanto o "L-GATr-slim" funciona melhor se você encolher a largura das camadas. Mesmo nesse tamanho microscópico, eles ainda superaram modelos de IA mais antigos que não eram conscientes da física.

3. A Versão "Quantizada" (Matemática de Baixa Precisão)

Este é o economizador de energia mais drástico. A IA padrão usa matemática muito precisa (como medir uma distância até o bilionésimo de milímetro). Os autores perceberam que, para o jet tagging, não é necessária tanta precisão. É possível obter bons resultados arredondando os números significativamente.

• A Analogia: Imagine que você está contando maçãs em um armazém.
  • IA Padrão: Você pesa cada uma das maçãs até o micrograma. (Preciso, mas leva muito tempo e consome muita energia da balança).
  • IA Quantizada: Você apenas as conta em números inteiros. (Rápido, consome quase nenhuma energia e, para o propósito de saber "quantas maçãs há", é perfeitamente adequado).
• O Método: Eles usaram uma técnica chamada PARQ (Quantização Regularizada de Afim por Partes). Pense nisso como uma regra de arredondamento inteligente que gentilmente empurra os números para serem simples (como 0, 1 ou -1) durante o processo de treinamento, em vez de forçá-los abruptamente.
• O Resultado: Ao mudar para esses números mais "brutos", eles reduziram o custo energético de execução da IA em 10 vezes (uma ordem de magnitude). A IA tornou-se incrivelmente rápida e eficiente energeticamente, com apenas uma pequena queda na precisão.

O Panorama Geral

Os autores combinaram essas três estratégias — Enxugar a arquitetura, Miniaturizar o tamanho e Quantizar a matemática — para criar "Jet Taggers Econômicos".

• Por que isso importa? Atualmente, esses modelos de IA poderosos são grandes demais para rodar no hardware que decide em tempo real quais colisões manter e quais descartar (o "trigger").
• O Objetivo: Ao tornar esses modelos pequenos, rápidos e eficientes energeticamente, os autores esperam, futuramente, rodá-los diretamente no hardware do trigger. Isso permitiria que o LHC usasse IA para tomar decisões de milésimos de segundo sobre quais colisões de partículas salvar, potencialmente descobrindo novas físicas que antes eram perdidas porque os dados eram descartados rápido demais.

Em resumo: Eles pegaram uma IA gigante e faminta por energia, colocaram-na em uma dieta, encolheram-na e ensinaram-na a fazer matemática com menos casas decimais, resultando em um motor minúsculo e super eficiente que ainda consegue reconhecer as partículas mais importantes do universo.

Resumo técnico

Resumo Técnico: Marcadores de Jatos Econômicos – Equivariantes, Enxutos e Quantizados

Problema
O aprendizado de máquina (ML) moderno transformou a marcação de jatos (jet tagging) no Grande Colisor de Hádrons (LHC), com transformadores Lorentz-equivariantes emergindo como arquiteturas de estado da arte. No entanto, modelos líderes como o L-GATr são computacionalmente caros, exigindo memória e tempo de treinamento significativos. Enquanto as tendências da indústria favorecem o escalonamento de redes e conjuntos de dados, a física do LHC enfrenta restrições específicas, particularmente em relação aos requisitos de memória e latência do hardware de gatilho (trigger). A classificação de jatos atual ainda não desempenha um papel no gatilho, mas os autores argumentam que deveria. O desafio central abordado é como reduzir o tamanho e o custo computacional dos modernos marcadores de jatos equivariantes, minimizando a degradação do desempenho, potencialmente permitindo sua implantação no nível do gatilho.

Metodologia
O artigo propõe uma estratégia de duas frentes para otimizar a eficiência de recursos: o enxugamento arquitetônico e a quantização numérica.

1. Arquitetura L-GATr-slim:
   Os autores introduzem uma versão simplificada do transformador Lorentz-equivariante (L-GATr). O L-GATr padrão utiliza uma representação de álgebra geométrica envolvendo escalares, pseudoescalares, vetores, axiais-vetores e tensores antissimétricos de posto dois. Os autores observam que pseudoescalares, axiais-vetores e tensores são desnecessários para a maioria das aplicações do LHC. Consequentemente, o L-GATr-slim restringe a representação latente apenas a escalares e vetores.

   • Camadas Lineares: Estendidas para operar em representações acopladas de escalar e vetor, garantindo que os componentes vetoriais compartilhem um único coeficiente escalar aprendível para manter a equivariância de Lorentz.
   • Não linearidade: Adapta a Unidade Linear de Portão (GLU) aplicando a não linearidade ao produto interno de dois vetores, multiplicado pelo vetor de saída.
   • Normalização: Modifica a RMSNorm para usar o valor absoluto do produto interno de Minkowski para canais vetoriais.
   • Atenção: Constrói matrizes de atenção escalares usando uma formulação específica que evita o produto externo computacionalmente caro usado no L-GATr completo.
   • Implementação: A arquitetura é projetada para ser compilada com torch.compile para maior eficiência.

2. Estratégias de Quantização:
   Os autores aplicam tipos de dados de baixa precisão e quantização de pesos para reduzir ainda mais os custos.

   • Quantização de Tipo de Dados: As entradas para as camadas lineares são quantizadas para int8 (usando quantização de ponto zero), mantendo o bfloat16 para operações sensíveis à precisão e para o passo de retropropagação (backward pass). Isso é aplicado às camadas ocultas do Transformer, ParT, L-GATr-slim e LLoCa-Transformer.
   • Quantização de Pesos: Os pesos lineares são quantizados para valores binários ou ternários usando Quantização de Gradiente Proximal (PARQ). Este método trata a quantização como uma restrição de regularização, utilizando um operador proximal para atualizar os pesos. Os autores comparam o PARQ com a Estimativa de Passagem Direta (STE), encontrando o PARQ com melhor estabilidade e desempenho.
   • Preservação da Equivariância: Um cuidado especial é tomado para garantir que a quantização não viole a equivariância de Lorentz. Para o LLoCa, a ortonormalização e as projeções de quadro permanecem em precisão total (float32), limitando as operações de baixa precisão aos invariantes de Lorentz. Para o L-GATr-slim, vetores completos são multiplicados por pesos quantizados, o que não introduz violações adicionais de simetria.

3. Escalonamento Ultra-Mini:
   Os autores investigam o desempenho dessas arquiteturas até 1.000 parâmetros, reduzindo o número de blocos ou a largura (canais) da rede.

Principais Resultados
O estudo compara os métodos propostos em três tarefas: marcação de topo (top tagging), regressão de amplitude e geração de eventos.

• Desempenho vs. Eficiência (L-GATr-slim):

  • No conjunto de dados JetClass (marcação de jatos multiclasse), o L-GATr-slim iguala o desempenho do L-GATr completo e do LLoCa-Transformer (AUC ~0,9885), mas reduz o tempo de treinamento por um fator de seis (de 166h para 27h em uma GPU H100) e o consumo de memória por um fator de dois.
  • Na regressão de amplitude ($Z + 4g$), o L-GATr-slim alcança o mesmo Erro Quadrático Médio (MSE) que o L-GATr completo, mas requer 20 vezes menos operações de treinamento e metade do tempo de treinamento.
  • Na geração de eventos ($t\bar{t} + nj$), a arquitetura enxuta iguala o desempenho de log-verossimilhança negativa dos modelos completos.

• Marcadores Ultra-Mini:

  • Ao reduzir o número de blocos (profundidade), o LLoCa-Transformer supera o L-GATr-slim em tamanhos muito pequenos (ex: 1.000 parâmetros).
  • Ao manter o número de blocos fixo (10) e reduzir os canais (largura), o L-GATr-slim mantém uma taxa de rejeição de fundo acima de 1.000 com apenas 2 canais vetoriais e 4 escalares, superando outras arquiteturas de 1.000 parâmetros.

• Ganhos de Quantização:

  • Quantizar as entradas para int8 e os pesos para valores ternários reduz o consumo de energia em aproximadamente uma ordem de magnitude (fator de 10) com perda de desempenho marginal.
  • O LLoCa-Transformer e o L-GATr-slim são robustos à quantização, mantendo alto desempenho onde transformadores padrão podem degradar mais significativamente.
  • Para o cenário mais restrito de recursos (1 bloco, espaço latente de 16 dimensões, int8), o LLoCa-Transformer quantizado (canonicidade global) retém um desempenho superior aos marcadores baseados em grafos prévios, apesar de uma redução de fator de dois na rejeição de fundo em comparação com sua contraparte de tamanho total.

Significância e Alegações
O artigo afirma que estas versões "econômicas" de transformadores equivariantes representam um caminho viável para a marcação de jatos no nível do gatilho no HL-LHC (High-Luminosity LHC). Ao combinar o enxugamento arquitetônico (remoção de componentes desnecessários de álgebra geométrica) e a quantização agressiva (PARQ e int8), os autores demonstram que é possível criar marcadores com ~1.000 parâmetros que mantêm os benefícios da equivariância de Lorentz motivados pela física.

Os autores enfatizam que, enquanto o escalonamento é o padrão da indústria, a física do LHC exige uma abordagem de "redução de escala consciente da física" (physics-aware downscaling). Os resultados sugerem que redes pequenas, quantizadas e equivariantes podem ser implantadas em hardware com recursos limitados (como FPGAs) sem sacrificar as simetrias fundamentais que tornam esses modelos eficazes, abrindo potencialmente novos caminhos para a análise em tempo real da subestrutura de jatos.