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

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

Planteamiento del Problema

Los Estados Cuánticos Neuronales (NQS) ofrecen un marco escalable para aproximar funciones de onda de muchos cuerpos cuánticos, particularmente cuando se utilizan modelos autorregresivos que permiten el muestreo exacto e independiente de la distribución de Born. Sin embargo, la optimización de estos modelos sigue siendo un cuello de botella significativo. Los métodos existentes enfrentan un compromiso drástico:

• Optimizadores de primer orden (p. ej., Adam): Computacionalmente eficientes y escalables, pero ignoran la geometría de la función de onda variacional, lo que a menudo conduce a una convergencia inestable o inexacta en aplicaciones de NQS.
• Reconfiguración Estocástica (SR) y sus variantes (p. ej., minSR): Geométricamente fundamentados (aproximando el descenso de gradiente natural), pero computacionalmente prohibitivos para redes grandes debido a la necesidad de resolver sistemas lineales grandes y, a menudo, mal condicionados. También sufren de fragilidad numérica en regímenes complejos.

Si bien los NQS autorregresivos resuelven los cuellos de botella de muestreo, la optimización estable y escalable para modelos de funciones de onda autorregresivas de gran tamaño sigue siendo un desafío abierto.

Metodología

Los autores proponen la Optimización de Funciones de Onda Proximal (PWO), un algoritmo de región de confianza inspirado en la Optimización de Política Próxima (PPO) del Aprendizaje por Refuerzo (RL). La metodología central se basa en una equivalencia formal entre la minimización de la energía variacional en NQS y la optimización de política de gradiente en RL.

1. Equivalencia RL-VMC: El artículo establece que, para Hamiltonianos estocásticos, el gradiente de la energía variacional puede expresarse como una actualización de gradiente de política ponderada por ventaja sobre la distribución de Born. En este mapeo:

   • Las configuraciones de espín corresponden a acciones.
   • Las energías locales centradas ($E_{loc}(s) - \mathbb{E}[E_{loc}]$) corresponden a ventajas.
   • La distribución de Born corresponde a la política.

2. Optimización de Doble Canal: A diferencia del PPO estándar que mane el una única distribución de probabilidad, las funciones de onda NQS son complejas, lo que requiere un manejo separado de la amplitud y la fase. PWO introduce dos objetivos de sustitución (surrogate) recortados:

   • Canal de Amplitud: Utiliza una razón de probabilidad recortada ($r_\theta = P_\theta / P_{\theta_{old}}$) para restringir las actualizaciones, imitando el PPO estándar. Esto asegura que la función de onda actualizada permanezca cerca de la referencia en términos de masa de probabilidad.
   • Canal de Fase: Introduce un nuevo sustituto recortado que restringe el incremento de fase envuelto ($\phi_\theta$) en lugar de una razón de probabilidad. Esto evita rotaciones abruptas de la función de onda mientras sigue la parte imaginaria del gradiente de Monte Carlo Variacional (VMC).

3. Implementación Algorítmica: PWO realiza múltiples actualizaciones internas ($K$) sobre un lote congelado de muestras extraídas de la distribución de referencia $P_{\theta_{old}}$. Reutiliza estas muestras a través de actualizaciones empleando el pesaje de importancia (con operadores de detención de gradiente para el canal de fase) para corregir los desajustes de la distribución. Esto evita la inversión de matrices explícita, manteniendo la escalabilidad de los métodos de primer orden mientras impone restricciones de región de confianza.

Contribuciones Clave

• Conexión Formal: Los autores conectan formalmente la minimización de la energía variacional con el gradiente de política de ventaja de RL, demostrando que el gradiente de NQS admite una forma ponderada por ventaja sobre la distribución de Born.
• Algoritmo PWO: Introducción de la Optimización de Función de Onda Proximal, un algoritmo de región de confianza que recorta los cambios en la razón de probabilidad en el canal de amplitud y los incrementos de fase en el canal de fase.
• Garantías Teóricas: Prueba de que el sustituto PWO satisface un límite de mejora de región de confianza y que la actualización de primer orden coincide con la dirección exacta del gradiente VMC cuando es on-policy.
• Demostración de Escalabilidad: Ajuste exitoso de un modelo de Gran Lenguaje (LLM) RWKV-7 de 1.5 mil millones de parámetros como un NQS autorregresivo, demostrando una optimización a una escala más de tres órdenes de magnitud por encima del trabajo previo.

Resultados Experimentales

Los autores evaluaron PWO frente a Adam, minSR y SPRING a través de varios benchmarks:

• Sistemas de Espín 1D (Ising y J1–J2):
  • Modelo de Ising: PWO convergió a un error relativo de $10^{-7}$ en aproximadamente 5 minutos, significativamente más rápido que minSR (~30 minutos) y Adam.
  • Cadena frustrada J1–J2: PWO mantuvo una convergencia estable con un V-score decreciente. En contraste, minSR fue numéricamente inestable (6 de cada 10 ejecuciones produjeron NaNs), y Adam se estancó en un error mucho mayor.
• Modelo J1–J2 2D Frustrado: En una red de $10 \times 10$, PWO alcanzó regímenes de energía más bajos más rápido que Adam y mantuvo una señal de error basada en la varianza (V-score) más baja dentro del mismo presupuesto de tiempo de ejecución.
• Escalabilidad: A medida que el tamaño del modelo aumentó (de NQS "pequeño" a "mediano"), PWO demostró estabilidad consistente y mejora de energía monotónica. Adam degradó su estabilidad y eficiencia para modelos más grandes, mientras que minSR no logró alcanzar valores de energía competitivos.
• Ajuste Fino a Gran Escala: El modelo RWKV-7 de 1.5B de parámetros ajustado con PWO logró errores relativos y V-scores finales más bajos que Adam, manteniendo una formación estable durante todo el presupuesto.

Significado y Reivindicaciones

El artículo afirma que PWO cierra la breancia entre la escalabilidad de los optimizadores de primer orden y la estabilidad de los métodos geométricos (como SR) mediante el aprovechamiento de los principios de región de confianza de RL.

• Estabilidad en Regímenes Difíciles: Los autores aseguran que los objetivos proximales tipo PPO proporcionan la estabilidad necesaria para entrenar NQS expresivos en Hamiltonianos cuánticos difíciles y frustrados donde los métodos tradicionales fallan o se vuelan inestables.
• Eficiencia de Muestreo: Al permitir la reutilización de muestras a través de múltiples actualizaciones sin la inversión de matrices explícita, PWO mejora la velocidad de convergencia en tiempo de ejecución.
• Escala: La aplicación exitosa a un modelo de 1.5B de parámetros sugiere que la optimización de NQS puede ahora entrar en el régimen de la modelación de secuencias modernas a gran escala. Los autores postulan que esto abre la puerta a investigar si la escala por sí sola puede superar las barreras actuales en la simulación de muchos cuerpos cuánticos, alineándose con la "lección amarga" (bitter lesson) de la IA.

El trabajo no pretende que los LLM sean el sesgo inductivo óptimo para Hamiltonianos 1D, sino que utiliza el modelo RWKV-7 para demostrar la eficacia del objetivo proximal a escalas sin precedentes.