OmniLearned: A Foundation Model Framework for All Tasks Involving Jet Physics

O artigo apresenta o OmniLearned, uma estrutura aprimorada de modelo fundamental para física de jatos que, ao combinar uma nova arquitetura, treinamento com mais de um bilhão de jatos e software documentado, estabelece o estado da arte em tarefas como identificação de quarks top e b e detecção de anomalias, ampliando o potencial de descoberta em experimentos de colisores.

O essencial

Imagine que você está tentando aprender a cozinhar. No passado, para aprender a fazer um prato específico (como um bolo de chocolate), você precisava comprar um livro de receitas só para bolo, outro só para salada, e outro só para carne. Cada livro tinha poucas receitas e você tinha que começar do zero cada vez que quisesse aprender algo novo. Isso era demorado e ineficiente.

Agora, imagine um Chef Mestre que passou anos comendo e estudando milhões de pratos diferentes: desde o básico até o mais sofisticado, de todas as cozinhas do mundo. Ele não apenas sabe fazer um bolo; ele entende a química dos ingredientes, como o calor age na massa e como os sabores se combinam.

Se você pedir para esse Chef fazer um prato novo (que ele nunca viu antes), ele não precisa de um livro de receitas. Ele usa todo o conhecimento geral que adquiriu para criar algo incrível rapidamente.

Este é o "OmniLearned" descrito no artigo.

Aqui está a explicação simplificada do que os cientistas fizeram:

1. O Problema: O "Chef" que precisava de mais ingredientes

Na física de partículas (especificamente no Large Hadron Collider - LHC), os cientistas estudam "jatos" (jets). Jatos são como "chuvas" de partículas que surgem quando colisões de alta energia acontecem. É difícil identificar o que causou cada chuva (se foi um quark, um glúon, ou algo novo e exótico).

Antes, os cientistas treinavam um modelo de Inteligência Artificial (IA) específico para cada tarefa. Se quisessem detectar um tipo específico de partícula, treinavam um modelo do zero. Se quisessem detectar outra, treinavam outro. Era como ter que estudar um novo livro de receitas para cada prato. Além disso, os dados de simulação eram caros e limitados.

2. A Solução: O "Chef Mestre" (OmniLearned)

Os autores criaram um novo modelo chamado OmniLearned. Pense nele como o Chef Mestre mencionado acima.

• O Treinamento Gigante: Em vez de usar 100 milhões de exemplos (como no modelo anterior), eles treinaram esse novo modelo com mais de 1 bilhão de jatos. É como se o Chef tivesse provado 10 vezes mais pratos do que qualquer outro antes dele.
• A Arquitetura (A Mente do Chef): Eles melhoraram a "mente" do modelo. Em vez de apenas olhar para os ingredientes, o modelo agora entende como os ingredientes se relacionam entre si (como a distância entre eles e como eles interagem). Eles usaram uma técnica chamada "Transformador" (a mesma usada em IAs de linguagem como o ChatGPT), mas adaptada para física.
• Flexibilidade: O modelo aprendeu a reconhecer padrões gerais. Se você der a ele um novo prato (uma nova tarefa), ele não precisa reaprender tudo. Você só precisa dar um "ajuste fino" (fine-tuning) rápido, e ele se adapta.

3. O Que Eles Provaram (Os Testes na Cozinha)

Para mostrar que o Chef é realmente bom, eles fizeram três testes principais:

• Teste 1: Identificar o "Rei" (Top Quark): Eles pediram para o modelo encontrar jatos que vinham de um "Top Quark" (uma partícula pesada e importante) em meio a milhões de outros jatos comuns. O OmniLearned foi o melhor de todos, superando os melhores modelos existentes.
• Teste 2: O "Detetive de Sabores" (B-Tagging): Na física, partículas têm "sabores" (como quark bottom, charm, etc.). O modelo foi capaz de distinguir esses sabores com muito mais precisão do que os métodos atuais usados pelo experimento ATLAS.
• Teste 3: Encontrar o "Intruso" (Detecção de Anomalias): Este é o mais legal. Imagine que você tem uma festa e quer encontrar alguém que não deveria estar lá, mas você não sabe quem é. O OmniLearned foi usado para olhar dados reais do experimento CMS e encontrar "anomalias". Ele conseguiu "redescobrir" o Top Quark (que já era conhecido, mas serviu como teste) apenas procurando por padrões estranhos nos dados, sem precisar que alguém dissesse exatamente o que procurar.

4. Por que isso é importante?

• Economia de Tempo e Dinheiro: Como o modelo já sabe "cozinhar" de tudo, os cientistas não precisam gastar meses treinando modelos do zero para cada novo experimento.
• Descobertas Futuras: O modelo é tão bom que pode ajudar a encontrar nova física. Se houver uma partícula nova e misteriosa no LHC, o OmniLearned pode ser a ferramenta que a identifica primeiro, porque ele sabe o que é "normal" e consegue notar o que é "estranho" com facilidade.
• Software Aberto: Eles não guardaram a receita. Eles liberaram o código e os dados para que qualquer cientista no mundo possa usar esse "Chef Mestre" em seus próprios experimentos.

Resumo em uma frase

Os cientistas criaram uma Super-IA treinada com mais de 1 bilhão de exemplos de colisões de partículas, que funciona como um "cérebro universal" capaz de identificar partículas, detectar erros e encontrar novas descobertas na física muito mais rápido e com mais precisão do que qualquer método anterior.

Resumo técnico

Resumo Técnico: OmniLearned – Um Framework de Modelo Fundamental para Física de Jatos

1. O Problema

Na física de colisores de alta energia, os jatos hadrônicos (resultantes de quarks e glúons de alta energia) são ubíquos e carregam informações cruciais sobre a origem das partículas e as propriedades da força forte. Tradicionalmente, a caracterização desses jatos exigia observáveis inspirados na física ou o treinamento de modelos de aprendizado de máquina específicos para cada tarefa (como identificação de top quarks, b-tagging ou detecção de anomalias).

O desafio central reside na dificuldade de reunir conjuntos de dados suficientemente grandes e variados para treinar modelos de ponta (state-of-the-art) para cada uma das inúmeras tarefas envolvendo jatos. Modelos treinados do zero (from scratch) para cada tarefa são computacionalmente custosos e não aproveitam o conhecimento adquirido em tarefas relacionadas. Embora modelos fundamentais (foundation models) baseados em aprendizado auto-supervisionado tenham sido explorados na física de partículas, muitas vezes não utilizam plenamente a continuidade dos dados ou as informações de rótulo disponíveis em simulações.

2. Metodologia e Arquitetura (OmniLearned)

O artigo apresenta o OmniLearned, uma atualização significativa do modelo anterior OmniLearn. A metodologia baseia-se em três pilares principais:

• Arquitetura PET v2 (Point Edge Transformer):

  • O modelo utiliza uma arquitetura Transformer combinada com uma Rede Neural de Grafos (GNN) chamada PET, adaptada para tratar jatos como nuvens de pontos com número variável de partículas.
  • Novas Features de Entrada: Além das informações cinemáticas básicas, o modelo incorpora diferenças em pseudorapidez ($\Delta\eta$) e ângulo azimutal ($\Delta\phi$), logaritmo do momento transversal ($\log p_T$) e energia ($\log E$).
  • Informação Adicional: O modelo lida de forma robusta com informações de identificação de partículas (PID) e vértices, utilizando blocos de embedding específicos que podem ser ativados ou desativados (com dropout durante o treinamento) para garantir generalização.
  • Atenção Local e Global: Foram introduzidos termos de interação inspirados na física nas camadas de atenção local e global. A atenção local usa um algoritmo de k-vizinhos mais próximos ($k=10$) com pesos aprendíveis, enquanto a atenção global incorpora termos de interação física na matriz de atenção.
  • Heads Específicos de Tarefa: O modelo possui cabeças para classificação e geração. Para classificação, utiliza-se cinco tokens aprendíveis (em vez de um) para resumir diferentes aspectos das partículas.

• Função de Perda e Treinamento Generativo:

  • O treinamento combina tarefas de classificação e geração.
  • Para a parte generativa, o modelo migrou de modelos de difusão tradicionais para Flow Matching. Isso permite prever o campo de velocidade ($u = \epsilon - x$) que conecta a distribuição de ruído aos dados, resultando em uma qualidade de geração superior e convergência mais rápida.
  • A função de perda inclui termos de entropia cruzada para classificação supervisionada e uma tarefa auxiliar de classificação de amostras (para dados sem rótulos de verdade), permitindo o uso de dados experimentais brutos no pré-treinamento.

• Conjunto de Dados (1 Bilhão de Jatos):

  • O modelo foi pré-treinado em um conjunto de dados massivo de 1,057 bilhões de jatos, agregando múltiplas fontes públicas (JetClass, JetClass2, Aspen Open Jets, ATLAS Top Tagging, H1 DIS, CMS QCD e BSM).
  • O conjunto abrange 210 classes (200 sabores de jatos específicos e 10 classes de conjuntos de dados), cobrindo diferentes sistemas de colisão (pp, ep) e simuladores (Delphes, Geant4).

3. Contribuições Principais

1. Framework OmniLearned: Uma arquitetura unificada capaz de aprender representações ricas de jatos que são facilmente adaptáveis a diversas tarefas downstream.
2. Escala de Dados: O uso de mais de 1 bilhão de jatos para pré-treinamento, aproveitando as propriedades de escalabilidade dos Transformers.
3. Software e Reprodutibilidade: Disponibilização de um pacote de software unificado para acesso a todos os conjuntos de dados, modelos pré-treinados (Small, Medium, Large) e código para fine-tuning.
4. Flexibilidade Arquitetural: Capacidade de incorporar informações adicionais (PID, vértices) e reutilizar cabeças de geração para tarefas de classificação (ex: previsão de origem de trajetórias).

4. Resultados

O desempenho do OmniLearned foi validado em três tarefas representativas:

• Classificação de Jatos (Top Tagging e Quark vs. Glúon):

  • No conjunto de dados de benchmark da comunidade para top tagging, o OmniLearned (versão Large) atingiu o estado da arte, superando modelos como ParT, L-GATr e o próprio OmniLearn anterior.
  • Para b-tagging (identificação de jatos contendo quarks bottom) usando simulações completas do ATLAS, o modelo superou a arquitetura GN2 (o padrão atual do ATLAS), melhorando a rejeição de jatos leves em mais de 50% e a rejeição de jatos $\tau$ em um fator de 2.
  • Observou-se uma forte correlação entre o tamanho do modelo e o desempenho após o fine-tuning.

• Detecção de Anomalias (CMS Open Data):

  • O modelo foi aplicado para redescobrir o top quark em dados reais do CMS (canal de jato único de alto $p_T$), onde os jatos de top representam apenas ~0,1% dos dados.
  • Estratégia 1 (Bump Hunt): Uso da capacidade generativa para estimar o fundo e da capacidade discriminativa para classificar anomalias. O OmniLearned detectou o sinal com significância estatística acima do limiar de descoberta, enquanto modelos treinados do zero falharam.
  • Estratégia 2 (Sem Fine-tuning): Uso direto das classes de pré-treinamento (razão entre decays de 3 pontas e QCD) para pontuação de anomalia, demonstrando que o modelo fundamental pode ser usado diretamente sem ajuste adicional.

5. Significado e Perspectivas

O OmniLearned representa um marco na aplicação de modelos fundamentais à física de colisores. Ele demonstra que:

• A pré-treinagem em grandes conjuntos de dados heterogêneos melhora drasticamente o desempenho em tarefas específicas, reduzindo a necessidade de simulações completas e custosas para cada novo experimento.
• A abordagem é generalizável entre diferentes detectores (ATLAS, CMS, H1) e sistemas de colisão.
• A metodologia pode ser estendida além da física de jatos, incluindo topologias completas de eventos e outros sistemas de colisão.

O trabalho estabelece um novo padrão para o desenvolvimento de algoritmos de aprendizado de máquina na física de partículas, focando na reutilização de representações aprendidas e na eficiência computacional, ao mesmo tempo que oferece ferramentas acessíveis à comunidade para acelerar a descoberta de nova física.