Vision Transformers and Graph Neural Networks for Charged Particle Tracking in the ATLAS Muon Spectrometer

Este artigo apresenta duas abordagens baseadas em aprendizado de máquina para otimizar o rastreamento de múons no Espectrômetro de Múons do ATLAS: o uso de Redes Neurais de Grafos para rejeitar ruídos e acelerar a reconstrução em 15%, e a aplicação de Vision Transformers para alcançar uma reconstrução aproximada ultra-rápida em 2,3 ms com 98% de eficiência, visando enfrentar os desafios de alta luminosidade do LHC.

O essencial

Imagine que o LHC (Grande Colisor de Hádrons) é uma máquina de fazer "explosões" de partículas subatômicas, e o experimento ATLAS é um detector gigante que tenta filmar o que acontece nessas explosões. O problema é que, no futuro, essas explosões vão acontecer com uma frequência e intensidade muito maiores do que hoje. Será como tentar encontrar uma agulha em um palheiro, mas o palheiro vai ficar 3 vezes maior e cheio de palhas falsas que parecem agulhas.

O foco deste trabalho é a parte do detector chamada Espectro de Muons, que é como uma "câmera de segurança" gigante que tenta rastrear partículas específicas (os múons) que escapam da explosão.

Aqui está a explicação simples das duas soluções inteligentes que os cientistas propuseram, usando analogias do dia a dia:

O Problema: O Palheiro Barulhento

Atualmente, quando os dados chegam, há um monte de "ruído" (partículas que não são importantes, como poeira cósmica ou erros eletrônicos). O computador tenta conectar os pontos (os "cliques" dos detectores) para desenhar a linha do caminho da partícula.

• O desafio: Com tanta gente (partículas) entrando ao mesmo tempo, o computador fica sobrecarregado, como um detetive tentando resolver 200 crimes ao mesmo tempo em uma sala cheia de gente gritando. Ele precisa ser mais rápido e mais esperto para ignorar o que não importa.

Solução 1: O "Filtro de Vizinhos" Inteligente (Redes Neurais de Grafos)

A primeira ideia é usar uma Rede Neural de Grafos (GNN).

• A Analogia: Imagine que você tem uma festa gigante onde as pessoas estão espalhadas. Você quer encontrar grupos de amigos que estão conversando (os múons), mas há muita gente gritando sozinha (o ruído).
• Como funciona: Em vez de olhar para cada pessoa individualmente, o algoritmo olha para "grupos" (agrupamentos de dados chamados "Buckets"). Ele pergunta: "Essas pessoas estão perto umas das outras e se comportam como um grupo?"
• O Resultado: O algoritmo consegue identificar e expulsar 97% das pessoas que estão apenas gritando sozinhas (o ruído) antes mesmo de tentar desenhar o caminho dos amigos.
• O Ganho: Ao limpar a festa antes de começar a contagem, o trabalho do detetive fica 15% mais rápido. É como se o detetive gastasse menos tempo limpando a sala e mais tempo resolvendo o caso.

Solução 2: O "Detetive com Óculos de Visão de Águia" (Vision Transformers)

A segunda ideia é muito mais ousada: tentar fazer todo o trabalho de rastreamento de uma só vez, usando uma tecnologia chamada Vision Transformer (ViT), que é a mesma usada em reconhecimento de imagens (como identificar gatos em fotos).

• A Analogia: Imagine que, em vez de tentar conectar ponto a ponto com uma régua (o método antigo), você joga uma foto de todo o caos da festa para um detetive com óculos de visão de águia. Esse detetive usa "atenção" (foco) para olhar apenas para onde importa.
• Como funciona:
  1. Filtro de Entrada: Primeiro, ele usa uma "lente" especial para ignorar instantaneamente 99,7% do ruído. De 6.900 pontos de dados, ele deixa apenas 55 pontos importantes. É como transformar uma sala cheia de gente gritando em uma sala com apenas 5 pessoas conversando.
  2. Rastreamento: Depois, ele usa uma técnica de "atenção em janelas" (Flash Attention). Imagine que ele olha para a foto em pedaços pequenos, mas consegue entender como os pedaços se conectam ao todo, sem precisar ler cada letra do livro. Ele desenha as linhas dos caminhos dos múons quase instantaneamente.
• O Resultado:
  • Velocidade: Ele consegue fazer isso em 2,3 milissegundos (menos de 3 milésimos de segundo) em uma placa de vídeo comum de computador (como as usadas por gamers).
  • Precisão: Ele acerta 98% dos múons corretamente.
  • O "Pulo do Gato": O método antigo demorava 255 milissegundos. O novo método é cerca de 100 vezes mais rápido em termos de tempo de processamento por evento.

Por que isso é importante?

No futuro, o LHC vai gerar dados tão rápido que o sistema atual não conseguirá acompanhar. Seria como tentar beber água de uma mangueira de incêndio com um canudinho.

Essas novas técnicas são como trocar o canudinho por um sistema de filtragem industrial:

1. Elas limpam a água (removem o ruído) antes de ela chegar ao copo.
2. Elas usam inteligência artificial para entender padrões complexos muito mais rápido do que os humanos ou softwares antigos conseguem.

Resumo final:
Os cientistas do ATLAS criaram dois "super-poderes" para o futuro:

1. Um filtro inteligente que limpa a bagunça antes de começar a trabalhar (GNN).
2. Um olho de águia que vê o caminho das partículas em frações de segundo, ignorando o caos ao redor (Vision Transformer).

Isso garante que, mesmo quando o LHC estiver no seu limite máximo de produção, a física ainda conseguirá descobrir novos segredos do universo sem ficar "travada" tentando processar dados inúteis.

Resumo técnico

1. Problema e Contexto

O experimento ATLAS no Grande Colisor de Hádrons (LHC) enfrenta desafios críticos com a chegada da era de Alta Luminosidade (HL-LHC), prevista para começar após 2030.

• Aumento de Pile-up: O número de colisões próton-próton por cruzamento de feixe (pile-up, $\langle\mu\rangle$) aumentará de ~60 (Run 3) para até 200.
• Desafio no Espectrômetro de Múons: Este aumento de densidade eleva a ocupação de dados no Espectrômetro de Múons, exigindo estratégias de processamento em tempo real mais eficientes e robustas, especialmente no Filtro de Eventos (Event Filter - EF) do sistema de gatilho (Trigger).
• Complexidade Específica: O espectrômetro possui uma baixa relação sinal-ruído (chegando a 0,6% em algumas regiões), integra cinco tecnologias de sub-detecção diferentes e opera em um campo magnético inhomogêneo.
• Limitação Atual: O algoritmo de reconstrução de base (baseline) atual, executado em CPU, leva cerca de 255 ms por evento para $\langle\mu\rangle=200$, o que é insuficiente para as taxas de entrada de 1 MHz do HL-LHC.

2. Metodologia

O trabalho propõe duas abordagens distintas baseadas em Aprendizado de Máquina (ML) para otimizar a reconstrução de múons:

Abordagem 1: Redes Neurais de Grafos (GNN) para Rejeição de Fundo

• Objetivo: Integrar-se à cadeia de reconstrução existente para rejeitar "hits" (sinais) de fundo antes da etapa de busca de trajetória (track finding).
• Estrutura do Grafos:
  • Os nós do grafo são "Buckets de Múons" (agrupamentos de hits dentro de 30 cm), e não hits individuais, para eficiência computacional.
  • A conectividade é baseada na proximidade espacial entre buckets em setores adjacentes.
• Arquitetura: Utiliza uma arquitetura baseada em EdgeConv (propagação de mensagens).
• Implementação: Integrada no framework Athena (C++), com desafios técnicos resolvidos na conversão de PyTorch Geometric para o formato ONNX para inferência.
• Função: Classificar buckets como contendo segmentos de múons, múltiplos segmentos ou apenas ruído de fundo.

Abordagem 2: Vision Transformers (ViT) para Reconstrução End-to-End

• Objetivo: Uma prova de conceito para realizar a reconstrução completa de múons (detecção, associação hit-rastreio e estimativa de parâmetros) puramente via ML.
• Arquitetura: Adaptação do modelo Mask2Former (originalmente para segmentação de imagens).
  • Tokenização: Trata cada hit individual do detector como um "token", evitando a necessidade de upsample complexo.
  • Atenção Janelada (Windowed Flash Attention): Implementa uma atenção global eficiente com complexidade $O(W \times N)$, ordenando os hits pelo ângulo azimutal ($\phi$) para respeitar a física (hits de um mesmo traço estão próximos em $\phi$).
  • Etapas:
    1. Filtro de Hits: Classificação binária para separar sinal de ruído.
    2. Rastreio (Tracking): Usa queries aprendidas para identificar candidatos a trajetórias e refinar limites iterativamente.
    3. Regressão: Estima parâmetros finais ($p_T$, $\eta$, $\phi$, carga).

3. Resultados Principais

Resultados do GNN (Rejeição de Fundo)

• Eficiência: Alcança uma taxa de rejeição de buckets de fundo de ~97% (em $\langle\mu\rangle=60$).
• Velocidade: Reduz o tempo de reconstrução por evento de 255 ms para 217 ms (uma melhoria de 15%) ao ser executado no pipeline Athena com GPU NVIDIA H100 e CPU AMD EPYC 9654.
• Qualidade Física: Preserva a eficiência e precisão da reconstrução de sinais em toda a faixa cinemática, sem introduzir viés negativo.

Resultados do ViT (Reconstrução End-to-End)

• Filtro de Hits:
  • AUC (Área sob a curva ROC) de 0.9997.
  • Aumenta a pureza de hits de 0,6% para 66,5% (rejeição de fundo de 99,7%) mantendo 99% de eficiência de sinal.
  • Reduz a ocupação média de 6.900 hits para apenas 55 hits por evento.
• Rastreio e Detecção:
  • Eficiência de Detecção de Sinal: 98,0% (para múons com $p_T \ge 5$ GeV e $\ge 3$ hits).
  • Taxa de Falsos Positivos: 5,1%.
  • Eficiência de Associação Hit-Rastreio: 92,9% (eficiência) e 88,90% (pureza).
  • Eficiência de Dupla Correspondência (Double Matching): 94,59% (onde tanto eficiência quanto pureza de associação são >50%).
  • Classificação de Carga: 96,35% de precisão.
• Velocidade de Inferência:
  • Em GPUs de nível consumidor (RTX 3090): 2,3 ms por evento (em lote de 200).
  • Em GPUs de alto desempenho (H100): 0,9 ms por evento.
  • Nota: A regressão de parâmetros é rápida (0,07 ms), mas o tempo total é limitado pela sobrecarga de lançamento de kernels da GPU.

4. Contribuições Chave

1. Integração Híbrida (GNN): Demonstração prática de como GNNs podem ser integrados em cadeias de reconstrução tradicionais (não-ML) para acelerar o processamento sem sacrificar a física, validado no framework de produção do ATLAS.
2. Reconstrução End-to-End com ViT: Primeira aplicação de arquiteturas Vision Transformers (Mask2Former adaptado) para o problema de rastreamento de partículas no Espectrômetro de Múons, tratando dados esparsos e heterogêneos como tokens de sequência.
3. Eficiência Computacional Extrema: Prova que é possível realizar a reconstrução de trajetórias complexas em milissegundos (2,3 ms) usando hardware acessível, viabilizando o processamento em tempo real para o HL-LHC.
4. Redução de Dados: O filtro de hits baseado em atenção reduz drasticamente a complexidade combinatória (de milhares para dezenas de hits), facilitando as etapas subsequentes.

5. Significado e Conclusão

Este trabalho valida o potencial de métodos baseados em ML para o pipeline de gatilho do ATLAS na era HL-LHC.

• Viabilidade: As abordagens demonstram que é possível lidar com a alta ocupação de dados ($\langle\mu\rangle=200$) mantendo a integridade física dos dados.
• Paralelismo: Embora a latência de um único evento no ViT pareça baixa, o ganho real vem da capacidade de processar lotes massivos em paralelo em GPUs, superando a limitação de threads de CPU.
• Futuro: O uso de mecanismos de atenção (Attention) é visto como uma escolha sustentável a longo prazo, dada a rápida evolução de hardware e software dedicado a essas arquiteturas.
• Próximos Passos: Otimização de tempo de execução (compilação de modelos, fusão de kernels), integração direta no Filtro de Eventos e adaptação para partículas de vida longa (LLP).

Em suma, o artigo oferece um caminho promissor para substituir ou complementar algoritmos tradicionais de rastreamento, garantindo que o experimento ATLAS possa operar com eficiência na próxima geração de colisões de alta energia.