Neural Quantum States in Mixed Precision

Resumo Técnico: Estados Quânticos Neurais em Precisão Mista

Definição do Problema
A computação científica tradicionalmente depende da aritmética de precisão dupla (64 bits) para garantir a precisão, particularmente em simulações de fenômenos do mundo real. No entanto, a crescente disponibilidade de aceleradores de hardware como GPUs tornou formatos de baixa precisão (ex: precisão simples f32, meia precisão f16, bfloat16) atraentes devido ao desempenho superior, redução da pegada de memória e melhor eficiência energética. Um gargalo significativo na computação científica moderna, especificamente nos métodos de Monte Carlo Variacional (VMC) para resolver sistemas de muitos corpos quânticos, é o custo computacional de gerar amostras de Monte Carlo de Markov (MCMC) suficientemente decorrelacionadas. Embora a aritmética de baixa precisão ofereça acelerações, seu impacto na precisão da amostragem MCMC — especificamente em relação à introdução de viés e ao potencial de destruição da ergodicidade — permanece subexplorado. Este trabalho aborda a necessidade de avaliar rigorosamente a aplicabilidade da aritmética de precisão mista em aplicações científicas baseadas em MCMC, com foco específico em Estados Quânticos Neurais (NQS).

Metodologia
Os autores abordam o problema por meio de uma combinação de derivação teórica e validação empírica:

1. Estrutura Teórica:

   • O estudo modela a precisão reduzida como uma perturbação aditiva $\delta(x)$ à log-probabilidade da distribuição alvo $\pi(x)$, resultando em uma distribuição perturbada $\tilde{\pi}(x)$.
   • Os autores derivam limites analíticos para o erro introduzido por essa perturbação no Metropolis-Hastings (MH) MCMC. Eles utilizam a Distância de Variação Total (TV) para quantificar a diferença entre as distribuições exata e perturbada.
   • Ao incorporar as propriedades de mistura da cadeia de Markov (especificamente usando a condição de minorização de Doeblin para estabelecer uma taxa de contração $r$), eles derivam um limite mais justo para o viés do que os limites grosseiros padrão.
   • Teorema III.3 fornece um limite específico para propostas de MH locais com ruído Gaussiano, mostrando que o viés escala com o nível de ruído $\sigma$ e a taxa de contração $r$. Crucialmente, se a cadeia mistura rapidamente ($r \ll 1$), o viés escala como $O(\sigma^2)$ em vez de $O(\sigma)$.

2. Validação Empírica:

   • Modelo de Brinquedo (Toy Model): Os autores primeiro validam os limites teóricos usando uma Máquina de Boltzmann Restrita (RBM) "ruidosa", onde o ruído Gaussiano é explicitamente injetado na log-densidade para simular erros de precisão.
   • Estrutura VMC de Precisão Mista: Eles implementam um esquema prático de precisão mista onde o componente de amostragem do VMC é executado em baixa precisão (f16 ou bf16), enquanto o cálculo de gradiente e as atualizações de parâmetros (passagens forward/backward) permanecem em precisão dupla (f64). Os parâmetros são atualizados em alta precisão e depois convertidos para baixo (downcast) para a etapa de amostragem.
   • Benchmarks: O método é testado no Modelo de Ising em Campo Transverso (TFIM) tanto nas fases antiferromagnética quanto paramagnética, bem como no modelo de Heisenberg. As arquiteturas testadas incluem RBMs e Redes Neurais Convolucionais Residuais (ResCNNs).

Principais Contribuições

1. Derivação de Limites de Erro: O artigo deriva limites analíticos gerais para o erro de amostragem introduzido pela aritmética de baixa precisão em MH MCMC discreto. Estabelece uma ligação teórica entre o nível de ruído, o tempo de mistura da cadeia e o viés resultante na distribuição estacionária.
2. Estrutura VMC de Precisão Mista: Os autores introduzem uma estrutura que realiza a amostragem MCMC em baixa precisão enquanto mantém alta precisão para o restante do loop de otimização. Isso efetivamente desacopla o gargalo de amostragem dos requisitos de precisão das atualizações de gradiente.
3. Demonstração Empírica de Escalabilidade: O trabalho demonstra que acelerações significativas (até 3,5$\times$) podem ser alcançadas na otimização de NQS sem degradar o desempenho do treinamento ou a precisão da aproximação do estado fundamental.

Resultados

• Concordância Teórica: Os resultados empíricos do modelo de brinquedo RBM ruidoso mostram excelente concordância com os limites teóricos derivados. O erro de amostragem de Monte Carlo nos valores esperados relativos permanece insensível a perturbações até que o parâmetro de ruído $\sigma$ seja da ordem de unidade.
• Preservação da Precisão: Nas otimizações completas de VMC para os modelos 2D TFIM e Heisenberg, o erro de energia relativa em relação a uma solução de referência de precisão dupla é quase idêntico entre os tipos de dados f64, f32, f16 e bf16.
• Análise de Ruído: A distribuição dos erros numéricos ($\delta$) nas avaliações de baixa precisão foi encontrada como aproximadamente Gaussiana, validando as suposições teóricas. O desvio padrão deste erro foi consistentemente pequeno ($\sigma < 1$), garantindo que as cadeias MCMC pudessem convergir de forma confiável.
• Tendências de Aceleração: Os fatores de aceleração aumentam sistematicamente com o número de amostras e o número de parâmetros do modelo. Isso é atribuído à melhor utilização da memória da GPU e à saturação das unidades aritméticas, respectivamente.
• Ruído de Gradiente: O estudo confirma que nos regimes f32 e f16, o ruído de amostragem estocástica de Monte Carlo domina o erro numérico de baixa precisão, explicando por que a dinâmica de otimização permanece estável.

Significância e Alegações
O artigo afirma que a aritmética de precisão mista é uma estratégia viável e eficaz para acelerar a amostragem MCMC na computação científica, especificamente para Estados Quânticos Neurais. Os autores afirmam que:

• Partes substanciais do algoritmo VMC, particularmente a amostragem do estado quântico, podem ser executadas em meia precisão sem perda de precisão.
• A estrutura teórica derivada fornece um método rigoroso para avaliar a aplicabilidade de precisão mista em abordagens de aprendizado de máquina que dependem de MCMC, estendendo-se além de NQS para áreas como aprendizado Bayesiano e modelos baseados em energia.
• Embora a precisão mista seja padrão no aprendizado de máquina geral, sua aplicação à computação científica baseada em MCMC tem sido subexplorada. Este trabalho fornece a justificativa teórica e a evidência prática necessárias para adotar essas técnicas para simulações mais escaláveis e energeticamente eficientes de sistemas quânticos de muitos corpos.
• Os autores observam que, embora os ganhos de amostragem sejam significativos, ganhos adicionais poderiam potencialmente ser realizados aplicando precisão mista aos cálculos de energia local e etapas de precondicionamento, embora isso apresente desafios práticos adicionais.

O trabalho conclui que a precisão mista é uma ferramenta de alto impacto para o futuro dos fluxos de trabalho científicos dependentes de GPU, desde que as propriedades de mistura das cadeias de Markov subjacentes sejam suficientemente robustas para lidar com o ruído numérico introduzido.