NetKet y JAX: cómo montar un modelo Transformer para sistemas de espín frustrados

Las arquitecturas Transformer empiezan a ganar espacio también en la física cuántica computacional. Una nueva guía práctica muestra cómo montar con NetKet y JAX un pipeline completo de Neural Quantum States para un problema complejo: la cadena de espín de Heisenberg J1-J2 frustrada. No es teoría en el vacío, sino un esquema reproducible en el que el modelo, el sampler, la optimización y la verificación de precisión forman un único circuito de investigación.

La idea central de la guía es que la arquitectura Transformer encaja bien con la descripción de estados cuánticos de muchos cuerpos, donde importan las correlaciones de largo alcance entre partículas. Los métodos numéricos habituales chocan pronto con la dimensión del espacio de estados, sobre todo si el sistema está frustrado, es decir, si interacciones en competencia le impiden alcanzar un mínimo de energía simple y ordenado. En estas condiciones, los Neural Quantum States permiten representar la función de onda como una red neuronal parametrizable y optimizarla después con Variational Monte Carlo. NetKet actúa como entorno listo para cálculos cuánticos, y JAX como motor de optimización de alta precisión y acelerada.

La guía configura primero 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 próximos interactúan con el coeficiente J1 y los siguientes con J2. Esa combinación crea la frustración que hace que el problema sea interesante y no trivial. Para describir el sistema se usan un grafo de NetKet, un espacio de Hilbert de espines 1/2 con proyección total fija y un operador hamiltoniano construido con GraphOperator. En paralelo se activa la precisión de 64 bits de JAX, importante para obtener cálculos estables en problemas de esta clase.

Después llega el aprendizaje automático propiamente dicho. La función de onda se define con un modelo propio, TransformerLogPsi, escrito en Flax: las configuraciones de espín se codifican como tokens, reciben embeddings y representaciones posicionales y pasan por varios bloques de self-attention y capas feed-forward. El ejemplo usa una dimensión oculta de 96, cuatro cabezas de atención y seis capas de transformer.

A la salida, 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 de forma adecuada solo con un escalar real. Tras agregar la información de toda la cadena mediante un promedio, la red obtiene una representación global de la configuración y puede expresar correlaciones más complejas que los ansätze locales.

Otro valor de la guía es que no se queda en la definición del modelo. Para el entrenamiento, el autor monta un circuito VMC completo: el sampler MetropolisExchange, el estado variacional MCState, el optimizador Adam y Stochastic Reconfiguration, el equivalente del natural gradient descent para estados cuánticos. La configuración del ejemplo usa 4096 muestras, descarta los estados iniciales de las cadenas y ejecuta unas 250 iteraciones de optimización. Esta pila no sirve solo para lograr una energía baja, sino también para controlar la convergencia. El código guarda las trayectorias de la energía media y de la varianza, para ver con qué estabilidad avanza el modelo hacia una buena solución.

Tras el entrenamiento, el pipeline se usa ya como herramienta de investigación. El autor ejecuta cálculos para varios valores de J2 entre 0 y 0,7 en una cadena de 24 nodos, registra las energías finales y estima el pico del factor de estructura. Así no se limita a ajustar los parámetros de la red: permite ver cómo cambia el comportamiento físico del sistema al aumentar la frustración y dónde pueden aparecer transiciones entre distintas fases de orden magnético.

Como comprobación adicional, se compara la calidad del modelo con la diagonalización exacta en un sistema menor, de 14 nodos, mediante el método de Lanczos. La comparación de energías ofrece una referencia numérica clara de cuánto se acerca el transformer variacional a la solución exacta allí donde el cálculo exacto aún es posible.

El valor práctico del material está en que cierra la brecha entre dos mundos: las arquitecturas modernas de aprendizaje profundo y los problemas reales de la física computacional. Para los ingenieros de ML es un buen ejemplo de cómo usar un transformer fuera del texto, las imágenes y los datos tabulares estándar. Para los físicos es una plantilla clara para pasar de la idea abstracta de los Neural Quantum States a un experimento reproducible con métricas, benchmarks y observables concretos. Y para quienes trabajan en la intersección de ambos campos, la guía ofrece una base ampliable: pasar a redes más grandes, añadir simetrías, estudiar el entrelazamiento o construir simulaciones temporales más complejas.

Qué significa: el enfoque Transformer se está convirtiendo poco a poco en una herramienta de trabajo no solo para las tareas clásicas de IA, sino también para modelar sistemas cuánticos, donde el coste del error es alto y los métodos exactos se agotan pronto. Si NetKet y JAX ya forman parte de tu stack, este material ofrece un punto de partida casi listo para experimentos de nivel investigador.