One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective

Resumo Técnico: One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective

Declaração do Problema

Os Estados Quânticos Neurais (NQS) oferecem um framework escalável para aproximar funções de onda de muitos corpos quânticos, particularmente ao utilizar modelos autorregressivos que permitem a amostragem exata e independente da distribuição de Born. No entanto, a otimização desses modelos continua sendo um gargalo significativo. Os métodos existentes enfrentam um compromisso agudo:

• Otimizadores de primeira ordem (ex: Adam): Computacionalmente eficientes e escaláveis, mas ignoram a geometria da função de onda variacional, levando frequentemente a uma convergência instável ou imprecisa em aplicações de NQS.
• Reconfiguração Estocástica (SR) e variantes (ex: minSR): Geometricamente fundamentados (aproximando o gradiente natural), mas computacionalmente proibitivos para redes grandes devido à necessidade de resolver sistemas lineares grandes e, muitas vezes, mal condicionados. Eles também sofrem de fragilidade numérica em regimes complexos.

Embora os NQS autorregressivos resolvam os gargalos de amostragem, a otimização estável e escalável para modelos de função de onda autorregressivos de grande escala permanece um desafio em aberto.

Metodologia

Os autores propõem a Otimização de Função de Onda Proximal (PWO), um algoritmo de região de confiança inspirado na Otimização de Política Proximal (PPO) de Aprendizado por Reforço (RL). A metodologia central baseia-se em uma equivalência formal entre a minimização de energia variacional em NQS e a otimização de gradiente de política em RL.

1. Equivalência RL-VMC: O artigo estabelece que, para Hamiltonianos estocásticos, o gradiente da energia variacional pode ser expresso como uma atualização de gradiente de política ponderada pela vantagem sobre a distribuição de Born. Nesta mapeamento:

   • Configurações de spin correspondem a ações.
   • Energias locais centradas ($E_{loc}(s) - \mathbb{E}[E_{loc}]$) correspondem a vantagens.
   • A distribuição de Born corresponde à política.

2. Otimização de Canal Duplo: Diferente do PPO padrão, que lida com uma única distribuição de probabilidade, as funções de onda NQS são complexas, exigindo o tratamento separado da amplitude e da fase. A PWO introduz dois objetivos substitutos (surrogates) com recorte (clipping):

   • Canal de Amplitude: Utiliza uma razão de probabilidade recortada ($r_\theta = P_\theta / P_{\theta_{old}}$) para restringir as atualizações, espelhando o PPO padrão. Isso garante que a função de onda atualizada permaneça próxima à de referência em termos de massa de probabilidade.
   • Canal de Fase: Introduz um novo substituto recortado que restringe o incremento de fase envolto ($\phi_\theta$) em vez de uma razão de probabilidade. Isso evita rotações abruptas da função de onda enquanto segue a parte imaginária do gradiente de Monte Carlo Variacional (VMC).

3. Implementação Algorítmica: A PWO realiza múltiplas atualizações internas ($K$) em um lote congelado de amostras extraídas da distribuição de referência $P_{\theta_{old}}$. Ela reutiliza essas amostras através de várias atualizações empregando ponderação por importância (com operadores de parada de gradiente para o canal de fase) para corrigir discrepâncias de distribuição. Isso evita a inversão explícita de matrizes, mantendo a escalabilidade dos métodos de primeira ordem enquanto impõe restrições de região de confiança.

Principais Contribuições

• Conexão Formal: Os autores conectam formalmente a minimização de energia variacional ao gradiente de política de vantagem de RL, demonstrando que o gradiente NQS admite uma forma ponderada pela vantagem sobre a distribuição de Born.
• Algoritmo PWO: Introdução da Otimização de Função de Onda Proximal, um algoritmo de região de confiança que recorta a razão de probabilidade no canal de amplitude e os incrementos de fase no canal de fase.
• Garantias Teóricas: Prova de que o substituto PWO satisfaz um limite de melhoria de região de confiança e que a atualização de primeira ordem coincide com a direção exata do gradiente VMC quando está no processo (on-policy).
• Demonstração de Escalabilidade: Ajuste bem-sucedido de um Modelo de Linguagem Grande (LLM) RWKV-7 de 1,5 bilhão de parâmetros como um NQS autorregressivo, demonstrando otimização em uma escala três ordens de magnitude superior ao trabalho anterior.

Resultados Experimentais

Os autores avaliaram a PWO contra Adam, minSR e SPRING através de várias métricas:

• Sistemas de Spin 1D (Ising e J1–J2):
  • Modelo de Ising: A PWO convergiu para um erro relativo de $10^{-7}$ em aproximadamente 5 minutos, significativamente mais rápido que o minSR (~30 minutos) e o Adam.
  • Cadeia Frustrada J1–J2: A PWO manteve uma convergência estável com um V-score decrescente. Em contraste, o minSR foi numericamente instável (6 de 10 execuções produziram NaNs), e o Adam estagnou em um erro muito mais alto.
• Modelo Frustrado 2D J1–J2: Em uma rede $10 \times 10$, a PWO atingiu regimes de energia mais baixos mais rapidamente que o Adam e manteve um sinal de erro baseado em variância (V-score) mais baixo dentro do mesmo orçamento de tempo de execução (wall-clock).
• Escalabilidade: À medida que o tamanho do modelo aumentava (de "tiny" para "medium" NQS), a PWO demonstrou estabilidade consistente e melhoria energética monotônica. O Adam degradou em estabilidade e eficiência para modelos maiores, enquanto o minSR falhou em atingir valores de energia competitivos.
• Ajuste Fino de Grande Escala: O modelo RWKV-7 de 1,5B de parâmetros, ajustado com PWO, alcançou erros relativos e V-scores finais menores que o Adam, mantendo o treinamento estável durante todo o orçamento.

Significância e Alegações

O artigo afirma que a PWO preenche a lacuna entre a escalabilidade dos otimizadores de primeira ordem e a estabilidade dos métodos geométricos (como SR) ao aproveitar os princípios de região de confiança do RL.

• Estabilidade em Regimes Difíceis: Os autores afirmam que os objetivos proximais do tipo PPO fornecem a estabilidade necessária para treinar NQS expressivos em Hamiltonianos quânticos difíceis e frustrados, onde métodos tradicionais falham ou tornam-se instáveis.
• Eficiência de Amostragem: Ao permitir a reutilização de amostras em múltiplas atualizações sem inversão explícita de matriz, a PWO melhora a velocidade de convergência de tempo de execução.
• Escala: A aplicação bem-sucedida a um modelo de 1,5B de parâmetros sugere que a otimização de NQS pode agora entrar no regime dos modelos de sequência modernos de grande escala. Os autores postulam que isso abre as portas para investigar se a escala sozinha pode superar as barreiras atuais na simulação de muitos corpos quânticos, alinhando-se com a "lição amarga" (bitter lesson) da IA.

O trabalho não afirma que os LLMs são o viés indutivo ideal para Hamiltonianos 1D, mas sim utiliza o modelo RWKV-7 para demonstrar a eficácia do objetivo proximal em escalas sem precedentes.