Event Generation with Parallel Langevin Sampling and Learned Stein Diagnostics

Resumo Técnico: Geração de Eventos com Amostragem de Langevin Paralela e Diagnósticos de Stein Aprendidos

Enunciado do Problema
A fenomenologia de colisores de precisão enfrenta um gargalo computacional significativo na geração de amostras de eventos de Monte Carlo (MC), particularmente para estados finais de alta multiplicidade. O custo dominante surge da amostragem da seção de choque diferencial de colisão próton-próton, onde as avaliações de elementos de matriz são dispendiosas. Os métodos padrão atuais, baseados amplamente no algoritmo Vegas com melhorias inspiradas na física, dependem de amostragem por rejeição. Essa abordagem sofre com eficiências de desponderamento (unweighting) baixas conforme a multiplicidade aumenta, frequentemente degradando em várias ordens de magnitude. Embora abordagens recentes de aprendizado de máquina (ex: fluxos de normalização) tenham melhorado a eficiência, elas permanecem fundamentalmente dependentes da amostragem por rejeição e suas eficiências associadas.

Metodologia
Os autores propõem um framework alternativo baseado em Cadeias de Markov de Monte Carlo (MCMC) paralelizadas usando o Algoritmo de Langevin Subamortecido (ULA), aumentado por um diagnóstico de convergência orientado por dados.

• Amostragem de Langevin Subamortecido Paralela: Em vez de executar uma única cadeia longa, o método evolui muitas cadeias de Markov independentes em paralelo em aceleradores de hardware. Para eliminar a autocorrelação dentro da cadeia, apenas o estado terminal de cada cadeia é retido como um evento desponderado. As cadeias utilizam uma dinâmica de Langevin subamortecido, retendo momento entre os passos e aplicando uma correção de Metropolis-Hastings (MH) após cada proposta para corrigir o viés de discretização.
• Dinâmica de Espaço Espelhado: Como o espaço de fase físico é limitado ($x \in [0, 1]^d$) enquanto o ULA é naturalmente formulado em $\mathbb{R}^d$, a evolução é realizada em um "espaço espelhado" usando uma transformação elemento a elemento (função inversa do erro). Isso preserva as propriedades de convergência do algoritmo não restrito.
• Discrepância de Stein Aprendida (LSD) como Diagnóstico: Para determinar quando o conjunto de cadeias atingiu a estacionalidade (ou seja, a distribuição alvo), os autores empregam uma Discrepância de Stein Aprendida. Diferente das discrepâncias de Stein baseadas em kernel que escalam quadraticamente com o tamanho da amostra ($n^2$), a LSD utiliza um crítico de rede neural para estimar a discrepância linearmente com o número de amostras. Isso permite o monitoramento eficiente da convergência.
• Tratamento de Fronteiras Afiadas: Seções de choque físicas envolvem cortes de espaço de fase abruptos. Para garantir que o operador de Stein permaneça válido na presença dessas fronteiras rígidas, uma função de fronteira escalar $h(x)$ é introduzida. Esta função anula-se na superfície do corte, garantindo que a contribuição da fronteira para o operador de Stein seja tratada corretamente.
• Inicialização por Surrogato: Para reduzir o número de avaliações caras de elementos de matriz exatos e de gradientes, um surrogato de rede neural é treinado em um conjunto limitado de pontos do espaço de fase. As cadeias são primeiro evoluídas usando este surrogato para acelerar o equilíbrio, sendo então alternadas para a densidade alvo exata por uma sequência mais curta de passos antes da terminação.

Principais Contribuições

1. Framework Algorítmico: A integração de cadeias ULA paralelas com uma discrepância de Stein aprendida fornece um método para gerar eventos estritamente desponderados sem depender de amostragem por rejeição.
2. Diagnóstico de Convergência: A aplicação da Discrepância de Stein Aprendida como um estimador prático e orientado por dados para o tempo de relaxação ($t_{rel}$), permitindo que o sistema determine quando parar de evoluir as cadeias e aceitar as amostras.
3. Aceleração por Surrogato: Demonstrar que surrogatos simples de redes neurais podem reduzir substancialmente o número de avaliações de alvo exato necessárias para o relaxamento, atuando como um "warm start" para as cadeias MCMC.
4. Tratamento de Fronteiras: Uma formulação específica do operador de Stein que contabiliza os cortes de espaço de fase agudos comuns na física de colisores.

Resultos
O método foi aplicado à geração de eventos de $u\bar{u} \to Z + ng$ em nível de árvore para multiplicidades de glúons $n = 0, 1, 2, 3$.

• Convergência: A Discrepância de Stein Aprendida rastreou com sucesso a convergência das cadeias. Para as multiplicidades estudadas, o relaxamento exigiu apenas um número modesto de passos ULA de alvo exato, com um crescimento apenas leve conforme a multiplicidade aumentava.
• Ganhos de Eficiência: O uso de inicialização por surrogato reduziu significativamente o número de avaliações de elementos de matriz exatos. Por exemplo, para $Z+3g$, o tempo de relaxação caiu de 115 passos (sem surrogato) para 55 passos (com surrogato).
• Validação: As amostras terminais das cadeias ULA reproduziram as distribuições de observáveis físicos (ex: momento transversal $p_T$, escala de resolução de jato $\sqrt{d_{01}}$ e rapidez) dentro das incertezas estatísticas quando comparadas contra estimativas independentes de Vegas.

Significância e Alegações
O artigo alega que o MCMC baseado em gradiente paralelo é uma alternativa praticamente viável à geração de eventos baseada em rejeição, particularmente em regimes onde a eficiência de desponderamento degrada severamente com a multiplicidade. Os autores enfatizam que seu método gera eventos estritamente desponderados, evitando os pesos residuais frequentemente encontrados na amostragem por rejeição quando os pesos máximos são limitados para manter a eficiência.

O overhead computacional do método é descrito como negligenciável comparado ao custo das avaliações de elementos de matriz. A abordagem é naturalmente paralela, tornando-a bem adequada para execução em GPU/TPU. Os autores observam que, embora o estudo atual foque em um canal partônico fixo, elementos de matriz de nível de árvore e um mapeamento de espaço de fase simples, a combinação de cadeias paralelas, diagnósticos explícitos de tempo de relaxação e inicialização por surrogato oferece uma rota concreta para a geração de eventos que minimiza as avaliações do alvo exato dispendioso. Eles permanecem modestos, afirmando que comparações mais amplas com os métodos atuais e o comportamento de escala devem ser testadas, em última análise, dentro de programas completos de geradores de eventos.