NetKet y JAX: cómo construir un modelo Transformer para sistemas de espines frustrados

Las arquitecturas Transformer están comenzando a encontrar su lugar en la física cuántica computacional: una nueva guía práctica demuestra cómo ensamblar un pipeline completo de Neural Quantum States usando NetKet y JAX para un problema complejo — la cadena de espines J1-J2 de Heisenberg frustrada. No es teoría abstracta, sino un marco reproducible donde el modelo, muestreador, optimización y verificación de precisión se reúnen en un único bucle de investigación. La idea central de la guía es que la arquitectura Transformer es muy adecuada para describir estados cuánticos de muchos cuerpos, donde las correlaciones de largo alcance entre partículas son importantes.

Los métodos numéricos convencionales rápidamente chocan contra la maldición de la dimensionalidad del espacio de estados, especialmente si el sistema está frustrado, es decir, las interacciones competitivas impiden que alcance un mínimo de energía ordenado simple. En tales condiciones, Neural Quantum States permiten que la función de onda se represente como una red neuronal parametrizada, que se optimiza luego mediante Monte Carlo Variacional. NetKet actúa como un entorno listo para computaciones cuánticas, mientras que JAX es el motor para optimización acelerada de alta precisión.

La guía primero configura la parte física básica del problema. El autor define una cadena unidimensional de longitud L con condiciones de contorno periódicas, donde los vecinos más cercanos interactúan con coeficiente J1, y sus vecinos siguientes con J2. Es precisamente esta combinación la que crea la frustración que hace que el problema sea interesante y no trivial.

Para describir el sistema, se utiliza un grafo en NetKet, un espacio de Hilbert de partículas de espín 1/2 con proyección total fija, y un operador Hamiltoniano ensamblado a través de GraphOperator. En paralelo, se activa la precisión de 64 bits en JAX, lo cual es esencial para cálculos estables en esta clase de problemas. Entonces comienza el aprendizaje automático en su forma pura.

La función de onda se define mediante un modelo TransformerLogPsi personalizado en Flax: las configuraciones de espín se codifican como tokens, reciben incrustaciones y representaciones posicionales, luego pasan a través de varios bloques de self-attention y capas feed-forward. El ejemplo utiliza una dimensión oculta de 96, cuatro cabezas de atención y seis capas Transformer. El modelo devuelve el logaritmo complejo de la amplitud de la función de onda — esto es crítico porque un estado cuántico no puede describirse adecuadamente solo con un escalar real.

Después de agregar información en toda la cadena mediante promediado, la red obtiene una representación global de la configuración y puede expresar correlaciones más complejas que los ansatze locales. Un valor particular de la guía es que no se detiene en la definición del modelo. Para el entrenamiento, el autor ensambla un bucle VMC completo: un muestreador MetropolisExchange, estado variacional MCState, optimizador Adam y Stochastic Reconfiguration como análogo del descenso de gradiente natural para estados cuánticos.

La configuración del ejemplo utiliza 4096 muestras, rechazo de estados iniciales en cadenas y aproximadamente 250 iteraciones de optimización. Tal stack es necesario no solo para lograr energía baja, sino también para controlar la convergencia. El código guarda trayectorias de energía promedio y varianza, para que se pueda ver cuán estable se mueve el modelo hacia una buena solución.

Después del entrenamiento, el pipeline se utiliza como herramienta de investigación. El autor ejecuta cálculos para varios valores de J2 en el rango de 0 a 0,7 para una cadena de longitud 24 nodos, registra las energías finales y estima el pico del factor de estructura. Esto permite no solo ajustar los parámetros de la red neuronal, sino observar cómo cambia el comportamiento físico del sistema cuando aumenta la frustración y dónde pueden aparecer transiciones entre diferentes fases de orden magnético.

Para verificación de calidad adicional, el modelo se compara con diagonalización exacta en un sistema más pequeño de tamaño 14 nodos mediante el método de Lanczos. La comparación de energías proporciona un criterio numérico claro: cuán cercano está el Transformer variacional a la solución exacta donde el cálculo exacto aún es viable. El significado práctico del material es que cierra la brecha entre dos mundos — arquitecturas modernas de aprendizaje profundo y problemas reales de física computacional.

Para ingenieros de ML, es un buen ejemplo de cómo un Transformer puede usarse fuera de texto, imágenes y datos tabulares estándar. Para físicos, es una plantilla comprensible para la transición de la idea abstracta de Neural Quantum States a un experimento reproducible con métricas concretas, benchmarks y cantidades observables. Y para aquellos que trabajan en la intersección de estos campos, la guía proporciona una base que puede extenderse: pasar a redes más grandes, añadir simetrías, estudiar entrelazamiento o construir simulaciones temporales más complejas.

Lo que esto significa: el enfoque Transformer se está convirtiendo gradualmente en una herramienta de trabajo no solo para tareas clásicas de IA, sino también para modelar sistemas cuánticos, donde el costo del error es alto y los métodos exactos se agotan rápidamente. Si NetKet y JAX ya están en su stack de trabajo, este material proporciona un punto de partida prácticamente listo para experimentos de nivel de investigación.