Neural Quantum States in Mixed Precision

Resumen Técnico: Estados Cuánticos Neuronales en Precisión Mixta

Planteamiento del Problema
La computación científica depende tradicionalmente de la aritmética de doble precisión (64 bits) para garantizar la exactitud, particularmente en simulaciones de fenómenos del mundo real. Sin embargo, la creciente disponibilidad de aceleradores de hardware como las GPU ha hecho que los formatos de baja precisión (por ejemplo, precisión simple f32, media f16, bfloat16) resulten atractivos debido a su rendimiento superior, menor huella de memoria y mejor eficiencia energética. Un cuello de botella significativo en la computación científica moderna, específicamente en los métodos de Monte Carlo Variacional (VMC) para resolver sistemas de muchos cuerpos cuánticos, es el costo computacional de generar muestras de cadenas de Markov Monte Carlo (MCMC) suficientemente descorrelacionadas. Si bien la aritmética de baja precisión ofrece aceleraciones, su impacto en la exactitud del muestreo MCMC —específicamente respecto a la introducción de sesgo y el potencial de destrucción de la ergodicidad— permanece insuficientemente explorado. Este trabajo aborda la necesidad de evaluar rigurosamente la aplicabilidad de la aritmética de precisión mixta en aplicaciones científicas impulsadas por MCMC, con un enfoque específico en los Estados Cuánticos Neuronales (NQS).

Metodología
Los autores abordan el problema mediante una combinación de derivación teórica y validación empírica:

1. Marco Teórico:

   • El estudio modela la reducción de precisión como una perturbación aditiva $\delta(x)$ a la log-probabilidad de la distribución objetivo $\pi(x)$, lo que resulta en una distribución perturbada $\tilde{\pi}(x)$.
   • Los autores derivan límites analíticos sobre el error introducido por esta perturbación en el algoritmo de Metropolis-Hastings (MH) MCMC. Utilizan la distancia de Variación Total (TV) para cuantificar la diferencia entre las distribuciones exacta y perturbada.
   • Al incorporar las propiedades de mezcla de la cadena de Markov (específicamente utilizando la condición de minoración de Doeblin para establecer una tasa de contracción $r$), derivan un límite más ajustado sobre el sesgo que los límites gruesos estándar.
   • El Teorema III.3 proporciona un límite específico para propuestas de MH locales con ruido Gaussiano, mostrando que el sesgo escala con el nivel de ruido $\sigma$ y la tasa de contracción $r$. Crucialmente, si la cadena se mezcla rápidamente ($r \ll 1$), el sesgo escala como $O(\sigma^2)$ en lugar de $O(\sigma)$.

2. Validación Empírica:

   • Modelo de Juguete: Los autores validan primero los límites teóricos utilizando una Máquina de Boltzmann Restringida (RBM) "con ruido", donde se inyecta explícitamente ruido Gaussiano en la log-densidad para simular errores de precisión.
   • Marco de VMC de Precisión Mixta: Implementan un esquema práctico de precisión mixta donde el componente de muestreo de VMC se ejecuta en baja precisión (f16 o bf16), mientras que el cálculo de gradientes y las actualizaciones de parámetros (pasadas hacia adelante/atrás) se mantienen en doble precisión (f64). Los parámetros se actualizan en alta precisión y luego se convierten a baja precisión para el paso de muestreo.
   • Benchmarks: El método se prueba en el Modelo de Ising en Campo Transverso (TFIM) tanto en fases antiferromagnéticas como paramagnéticas, así como en el modelo de Heisenberg. Las arquitecturas probadas incluyen RBMs y Redes Neuronales Convolucionales Residuales (ResCNNs).

Contribuciones Clave

1. Derivación de Límites de Error: El artículo deriva límites analíticos generales sobre el error de muestreo introducido por la aritmética de baja precisión en MCMC de MH discreto. Establece un vínculo teórico entre el nivel de ruido, el tiempo de mezcla de la cadena y el sesgo resultante en la distribución estacionaria.
2. Marco de VMC de Precisión Mixta: Los autores introducen un marco que realiza el muestreo MCMC en baja precisión manteniendo la alta precisión para el resto del bucle de optimización. Esto desacopla efectivamente el cuello de botella del muestreo de los requisitos de precisión de las actualizaciones de gradientes.
3. Demostración Empírica de Escalabilidad: El trabajo demuestra que se pueden lograr aceleraciones significativas (hasta 3.5$\times$) en la optimización de NQS sin degradar el rendimiento del entrenamiento ni la exactitud de la aproximación del estado fundamental.

Resultados

• Acuerdo Teórico: Los resultados empíricos del modelo de juguete RBM con ruido muestran un excelente acuerdo con los límites teóricos derivados. El error de muestreo de Monte Carlo en los valores esperados relativos permanece insensible a las perturbaciones hasta que el parámetro de ruido $\sigma$ es del orden de la unidad.
• Preservación de la Exactitud: En las optimizaciones completas de VMC para los modelos 2D de TFIM y Heisenberg, el error de energía relativa con respecto a una solución de referencia de doble precisión es casi idéntico a través de los tipos de datos f64, f32, f16 y bf16.
• Análisis de Ruido: Se encontró que la distribución de los errores numéricos ($\delta$) en las evaluaciones de baja precisión es aproximadamente Gaussiana, validando los supuestos teóricos. La desviación estándar de este error fue consistentemente pequeña ($\sigma < 1$), asegurando que las cadenas MCMC pudieran converger de manera confiable.
• Tendencias de Aceleración: Los factores de aceleración aumentan sistemáticamente tanto con el número de muestras como con el número de parámetros del modelo. Esto se atribuye a una mejor utilización de la memoria de la GPU y a la saturación de las unidades aritméticas, respectivamente.
• Ruido de Gradiente: El estudio confirma que en los regímenes f32 y f16, el ruido de muestreo estocástico de Monte Carlo domina el error numérico de baja precisión, lo que explica por qué la dinámica de optimización permanece estable.

Significancia y Reivindicaciones
El artículo afirma que la aritmética de precisión mixta es una estrategia viable y efectiva para acelerar el muestreo MCMC en la computación científica, específicamente para los Estados Cuánticos Neuronales. Los autores sostienen que:

• Gran parte del algoritmo VMC, particularmente el muestreo del estado cuántico, puede ejecutarse en media precisión sin pérdida de exactitud.
• El marco teórico derivado proporciona un método riguroso para evaluar la aplicabilidad de la precisión mixta en enfoques de aprendizaje automático que dependen de MCMC, extendiéndose más allá de NQS hacia áreas como el aprendizaje Bayesiano y los modelos basados en energía.
• Si bien la precisión mixta es estándar en el aprendizaje automático general, su aplicación a la computación científica impulsada por MCMC ha sido insuficientemente explorada. Este trabajo proporciona la justificación teórica y la evidencia práctica necesarias para adoptar estas técnicas para simulaciones más escalables y energéticamente eficientes de sistemas cuánticos de muchos cuerpos.
• Los autores señalan que, aunque las aceleraciones de muestreo son significativas, se podrían obtener mayores ganancias aplicando la precisión mixta a los cálculos de energía local y a los pasos de precondicionamiento, aunque esto presenta desafíos prácticos adicionales.

El trabajo concluye que la precisión mixta es una herramienta de alto impacto para el futuro de los flujos de trabajo dependientes de GPU, siempre que las propiedades de mezcla de las cadenas de Markov subyacentes sean lo suficientemente robustas para manejar el ruido numérico introducido.