JEDI-linear: Fast and Efficient Graph Neural Networks for Jet Tagging on FPGAs

Este artigo apresenta o JEDI-linear, uma nova arquitetura de Rede Neural em Grafos de complexidade linear que utiliza quantização de alta granularidade e operações sem multiplicadores para alcançar latência e eficiência de recursos recordes em FPGAs, atendendo com sucesso aos requisitos de identificação de jatos em tempo real para o sistema de gatilho de Nível 1 do CMS no HL-LHC.

O essencial

Imagine que você está operando um posto de controle de segurança de alta velocidade em um aeroporto massivo (o Grande Colisor de Hádrons). A cada 25 nanossegundos, um novo "voo" de partículas colide com o solo, criando um spray caótico de detritos. Sua função é olhar instantaneamente para esse spray e decidir: "Isso é uma pilha entediante de lixo ou é um tesouro raro e valioso?"

Se você tentar salvar cada pedaço individual de detrito, ficará sem espaço de armazenamento em uma fração de segundo. Portanto, você precisa de um sistema de gatilho — um filtro ultra-rápido que toma decisões em frações de segundo para reter apenas os eventos interessantes.

É aqui que o artigo entra. Os autores construíram um novo "cérebro" ultra-rápido (chamado JEDI-linear) para ajudar esses guardas de segurança a tomar melhores decisões, e conseguiram integrar esse cérebro em um pequeno chip de computador especializado (um FPGA) que precisa funcionar incrivelmente rápido.

Abaixo está a explicação detalhada de sua invenção usando analogias simples:

1. O Problema: O Gargalo do "Aperto de Mão"

Métodos anteriores para classificar esses sprays de partículas (chamados de "jatos") usavam uma técnica semelhante a um aperto de mão rotativo massivo.

• O Jeito Antigo: Imagine uma sala com 64 pessoas. Para entender o grupo, o método antigo exigia que cada pessoa individualmente se virasse e apertasse a mão de todas as outras pessoas individualmente.
• O Resultado: Se você tem 64 pessoas, isso são mais de 4.000 apertos de mão. Leva tempo demais, e a sala fica muito lotada de pessoas tentando falar ao mesmo tempo. No mundo da física de partículas, esse processo de "aperto de mão" é lento demais e ocupa muito espaço de hardware para ser útil em verificações de segurança em tempo real.

2. A Solução: O "Reunião em Grupo" (JEDI-linear)

Os autores perceberam que não precisavam que todos apertassem as mãos individualmente. Em vez disso, inventaram uma abordagem de complexidade linear.

• O Novo Jeito: Em vez de apertos de mão individuais, imagine que todas as pessoas na sala simplesmente levantam a mão para compartilhar seu humor atual, e um único "capitão" reúne todos esses humores em um único grande resumo. Então, o capitão diz a todos: "Aqui está o clima de todo o grupo."
• A Magia: Agora, em vez de 4.000 apertos de mão, você precisa apenas que 64 pessoas falem uma vez. O trabalho escala linearmente (se você dobrar o número de pessoas, você dobra o trabalho, não o quadruplica). Esta é a parte "JEDI-linear": mantém o contexto do grupo sem as interações pareadas bagunçadas e lentas.

3. As Truques de Hardware: Fazendo Caber em um Chip Minúsculo

Mesmo com o novo método de "reunião", o cérebro ainda precisava ser pequeno e rápido o suficiente para caber em um tipo específico de chip usado no sistema de segurança. Os autores usaram dois truques inteligentes:

• O Truque do "Uniforme Personalizado" (Quantização):
  Geralmente, computadores tratam todos os números da mesma maneira (como dar a cada soldado o mesmo casaco pesado). Os autores perceberam que algumas partes da matemática são muito sensíveis e precisam de alta precisão (um casaco pesado), enquanto outras não se importam muito (uma camiseta leve). Eles treinaram o sistema para usar um "uniforme personalizado", atribuindo larguras de bits pequenas e eficientes a números que não precisam de muita precisão. Isso reduziu significativamente a pegada de memória.

• O Truque do "Sem Multiplicador" (Aritmética Distribuída):
  Chips padrão usam blocos especiais e caros de "multiplicadores" para fazer matemática, que são como motores pesados e famintos por energia. Os autores substituíram esses motores por um sistema inteligente de somadores e deslocadores (como usar uma régua de cálculo ou uma pilha de blocos).

  • O Resultado: Eles eliminaram completamente a necessidade dos pesados "motores multiplicadores" (blocos DSP). Isso economizou enormes quantidades de espaço e energia, permitindo que o sistema rodasse em um chip que anteriormente não conseguia suportar a carga.

4. Os Resultados: Velocidade e Eficiência

Quando testaram esse novo sistema contra os melhores métodos existentes:

• Velocidade: É 3,7 a 11,5 vezes mais rápido. Pode tomar uma decisão em menos de 60 nanossegundos (o que é mais rápido que um piscar de olhos).
• Eficiência: Usa até 150 vezes menos "tempo de inicialização" entre decisões e ocupa 6,2 vezes menos espaço no chip.
• Precisão: Apesar de ser menor e mais rápido, é na verdade mais preciso na identificação dos raros jatos de partículas do que os modelos anteriores, mais pesados.

Por Que Isso Importa

Os autores afirmam que esta é a primeira vez que um modelo de IA baseado em interações foi rápido e pequeno o suficiente para ser usado no sistema de Gatilho de Nível 1 no Grande Colisor de Hádrons de Alta Luminosidade do CERN.

Pense nisso como atualizar a segurança do aeroporto de uma busca manual e lenta para um scanner super-rápido e automatizado que nunca perde um item raro, mas nunca atrasa a fila. Isso permite que os cientistas capturem eventos físicos raros que anteriormente eram rápidos demais para serem vistos, tudo isso usando menos hardware do que uma calculadora padrão.

Em resumo: Eles pegaram uma IA complexa e lenta, simplificaram sua matemática para que ela não precise "conversar consigo mesma" constantemente, vestiram-na com roupas sob medida para economizar espaço e substituíram seus motores pesados por engrenagens leves. O resultado é um cérebro minúsculo e super-rápido que cabe em um chip e pode detectar partículas raras em tempo real.

Resumo técnico

1. Declaração do Problema

Experimentos de física de altas energias no Grande Colisor de Hádrons (LHC) do CERN geram volumes massivos de dados (centenas de terabytes por segundo). Para gerenciar isso, o sistema Level-1 Trigger (L1T) deve filtrar eventos em tempo real (dentro de alguns microssegundos) utilizando FPGAs.

• O Desafio: A identificação de jatos (identificação da origem de sprays de partículas) é crítica para essa filtragem. Embora Redes Neurais em Grafos (GNNs), especificamente Redes de Interação (INs) como JEDI-net, ofereçam superioridade em precisão ao modelar interações de partículas, elas são difíceis de implantar em FPGAs para o L1T devido a:
  • Complexidade Computacional: GNNs padrão exigem cálculos explícitos de arestas entre pares ($O(N^2)$), criando um gargalo para jatos com muitas partículas.
  • Restrições de Hardware: Requisitos estritos de latência (<100 ns), recursos limitados de FPGA (tipicamente <1 Super Logic Region) e a necessidade de intervalos de iniciação (II) baixos.
  • Consumo de Recursos: GNNs baseadas em FPGA existentes frequentemente exigem milhares de blocos de Processamento de Sinal Digital (DSP) e grandes contagens de Tabelas de Pesquisa (LUT), tornando-as impraticáveis para implantação no mundo real ao lado de outros algoritmos.

2. Metodologia

Os autores propõem o JEDI-linear, uma arquitetura GNN inovadora projetada especificamente para eficiência de hardware, combinada com técnicas avançadas de otimização.

A. Inovação Algorítmica: Complexidade Linear

• Conceito Central: O JEDI-net original computa interações entre cada par de partículas ($O(N^2)$). O JEDI-linear reformula a função de interação de aresta $f_R$ como uma simples transformação afim (uma única camada densa).
• Derivação Matemática: Ao assumir $f_R(I_i \| I_j) = W_1 I_i + W_2 I_j + C$, a soma explícita de pares pode ser reescrita como uma agregação global. O embedding de interação para a partícula $i$ torna-se uma função da média global de todos os recursos das partículas mais uma transformação dos recursos da partícula individual.
• Resultado: Isso reduz a complexidade computacional de quadrática $O(N^2)$ para linear $O(N)$, eliminando a necessidade de cálculos explícitos no nível de arestas, enquanto preserva o contexto global.

B. Estratégias de Otimização de Hardware

1. Treinamento Consciente de Quantização de Alta Granularidade (QAT):

   • Diferentemente da quantização uniforme, os autores utilizam uma abordagem de otimização de largura de bits por parâmetro.
   • Usando um gradiente substituto diferenciável, o processo de treinamento atribui automaticamente larguras de bits específicas a cada peso com base em seu impacto na precisão e no custo de hardware (medido por Operações de Bits Efetivas, EBOPs).
   • Isso permite modelos de precisão mista onde muitos pesos são podados (largura de bits reduzida a zero) ou reduzidos a 1-2 bits, encolhendo significativamente o tamanho do modelo sem sacrificar a precisão.

2. Aritmética Distribuída (DA) para MACs sem Multiplicadores:

   • Para reduzir ainda mais o uso de recursos, a implementação substitui multiplicadores convencionais por Aritmética Distribuída.
   • A DA decompõe multiplicações de matriz-vetor em operações de deslocamento-soma implementadas via LUTs.
   • Resultado: O design elimina completamente a necessidade de blocos DSP, confiando apenas em LUTs e registradores, que são mais abundantes e flexíveis em FPGAs.

3. Arquitetura de Fluxo de Dados Totalmente Desenrolada:

   • O design emprega um fluxo de dados estático e totalmente desenrolado, onde cada operação é mapeada para hardware dedicado.
   • Isso evita o compartilhamento de recursos e sobrecarga de controle, permitindo um intervalo de iniciação de 1 ciclo e latência ultra-baixa e determinística.

3. Principais Contribuições

• Arquitetura JEDI-linear: A primeira GNN baseada em interação para identificação de jatos que alcança complexidade linear ao remover interações explícitas entre pares, tornando-a escalável para grandes contagens de partículas.
• Co-Design Consciente de Hardware: Integração de quantização de precisão mista de alta granularidade e Aritmética Distribuída para criar uma implementação sem multiplicadores e sem DSP.
• Framework de Automação: Um framework da4ml estendido que rastreia automaticamente gráficos de computação simbólica e gera Verilog sintetizável para essas arquiteturas complexas e desenroladas.
• Código Aberto: Lançamento de templates e código do JEDI-linear para apoiar a reprodutibilidade.

4. Resultados Experimentais

Os modelos foram avaliados em FPGAs AMD VU13P visando o sistema de Trigger Level-1 do CMS (Camada de Correlador 2).

• Latência e Vazão:
  • Alcançou latência <60 ns (por exemplo, 52 ns para 16 partículas com 16 recursos) e intervalo de iniciação de 1 ciclo de relógio.
  • Isso representa uma latência 3,7x a 11,5x menor e um intervalo de iniciação de até 150x menor em comparação com designs de GNN de última geração (SOTA) (por exemplo, LL-GNN, variantes JEDI-net).
• Eficiência de Recursos:
  • 0 blocos DSP utilizados em todas as configurações (designs SOTA frequentemente usam 5.000–9.000+ DSPs).
  • Até 6,2x menor uso de LUT em comparação com modelos SOTA.
  • Exemplo: Um modelo JEDI-linear de 32 partículas usa 6,2x menos LUTs e alcança 11,5x menor latência do que o modelo GNN J5, oferecendo ao mesmo tempo maior precisão (81,4% vs 79,9%).
• Precisão:
  • Alcançou até 82,4% de precisão de classificação (em entradas de 16 recursos com 64 partículas).
  • Supera o DeepSets (DS) e implementações anteriores de GNN em várias contagens de partículas (8 a 128).
  • Demonstrou escalabilidade superior, mantendo alta precisão à medida que a contagem de partículas aumenta, enquanto outros modelos degradam-se ou tornam-se inviáveis.

5. Significado

• Primeira Viabilidade no Mundo Real: Esta é a primeira GNN para identificação de jatos a atender às restrições estritas de latência (<60 ns) e recursos do Trigger Level-1 do CMS HL-LHC. Permite o uso de GNNs poderosas em triggers de hardware em tempo real, uma tarefa anteriormente dominada por modelos mais simples e menos precisos.
• Escalabilidade: A complexidade linear permite que o sistema lide com jatos de grandes números de partículas (até 128) sem crescimento exponencial de recursos, garantindo a viabilidade futura do sistema para corridas de maior luminosidade.
• Impacto Mais Amplo: As técnicas (processamento linearizado, quantização de precisão mista, aritmética distribuída e designs totalmente desenrolados) são aplicáveis além da física de partículas a outros domínios de baixa latência, como DNNs confiáveis, VAEs e Transformers.

Em conclusão, o JEDI-linear conecta com sucesso a lacuna entre algoritmos de aprendizado profundo de alta precisão e as restrições extremas de triggers de hardware em tempo real, pavimentando o caminho para sistemas de trigger inteligentes de próxima geração no LHC.