NetKet et JAX : construire un modèle Transformer pour les systèmes de spins frustrés

Les architectures Transformer commencent aussi à s'imposer en physique quantique computationnelle. Un nouveau guide pratique montre comment construire avec NetKet et JAX un pipeline complet de Neural Quantum States pour un problème difficile : la chaîne de spins de Heisenberg J1-J2 frustrée. Il ne s'agit pas de théorie hors sol, mais d'un schéma reproductible où le modèle, l'échantillonneur, l'optimisation et la vérification de la précision forment une seule boucle de recherche.

L'idée centrale du guide est que l'architecture Transformer convient bien à la description d'états quantiques à N corps, où les corrélations à longue portée entre particules comptent. Les méthodes numériques classiques se heurtent vite à la dimension de l'espace des états, surtout si le système est frustré, c'est-à-dire si des interactions concurrentes l'empêchent d'atteindre un minimum d'énergie simple et ordonné. Dans ces conditions, les Neural Quantum States permettent de représenter la fonction d'onde sous forme de réseau de neurones paramétrable, puis de l'optimiser par Variational Monte Carlo. NetKet joue ici le rôle d'environnement prêt à l'emploi pour le calcul quantique, et JAX celui de moteur d'optimisation accélérée et de haute précision.

Le guide commence par régler la partie physique de base. L'auteur définit une chaîne unidimensionnelle de longueur L avec des conditions aux limites périodiques, où les plus proches voisins interagissent avec le coefficient J1 et les voisins suivants avec J2. C'est cette combinaison qui crée la frustration, et rend le problème intéressant et non trivial. Le système est décrit par un graphe NetKet, un espace de Hilbert de spins 1/2 à projection totale fixée et un opérateur hamiltonien construit via GraphOperator. La précision 64 bits de JAX est activée en parallèle, ce qui compte pour obtenir des calculs stables dans ce type de problème.

Vient ensuite l'apprentissage automatique proprement dit. La fonction d'onde est définie par un modèle maison, TransformerLogPsi, écrit en Flax : les configurations de spins sont encodées en tokens, reçoivent des embeddings et des représentations positionnelles, puis traversent plusieurs blocs de self-attention et de couches feed-forward. L'exemple utilise une dimension cachée de 96, quatre têtes d'attention et six couches de transformer.

En sortie, le modèle renvoie le logarithme complexe de l'amplitude de la fonction d'onde. C'est crucial, car un état quantique ne peut pas être décrit correctement par un simple scalaire réel. Après agrégation de l'information sur toute la chaîne par moyenne, le réseau obtient une représentation globale de la configuration et peut exprimer des corrélations plus complexes que les ansätze locaux.

La valeur du guide tient aussi au fait qu'il ne s'arrête pas à la définition du modèle. Pour l'entraînement, l'auteur assemble une boucle VMC complète : l'échantillonneur MetropolisExchange, l'état variationnel MCState, l'optimiseur Adam et la Stochastic Reconfiguration, équivalent de la descente de gradient naturelle pour les états quantiques. La configuration de l'exemple utilise 4096 échantillons, écarte les états initiaux des chaînes et effectue environ 250 itérations d'optimisation. Cette pile sert non seulement à obtenir une énergie basse, mais aussi à contrôler la convergence. Le code enregistre les trajectoires de l'énergie moyenne et de la variance, afin de voir avec quelle stabilité le modèle progresse vers une bonne solution.

Une fois l'entraînement terminé, le pipeline sert d'outil de recherche. L'auteur lance des calculs pour plusieurs valeurs de J2 entre 0 et 0,7 sur une chaîne de 24 sites, relève les énergies finales et estime le pic du facteur de structure. Il ne s'agit donc pas seulement d'ajuster les paramètres du réseau : on observe comment le comportement physique du système évolue quand la frustration augmente, et où des transitions entre différentes phases d'ordre magnétique peuvent apparaître.

Pour une vérification supplémentaire, la qualité du modèle est comparée à la diagonalisation exacte sur un système plus petit, de 14 sites, par la méthode de Lanczos. La comparaison des énergies fournit un repère numérique clair : elle montre à quel point le transformer variationnel s'approche de la solution exacte là où le calcul exact reste possible.

L'intérêt pratique de ce travail est de combler le fossé entre deux mondes : les architectures modernes de deep learning et les problèmes réels de la physique computationnelle. Pour les ingénieurs ML, c'est un bon exemple d'usage du transformer hors du texte, des images et des données tabulaires standard. Pour les physiciens, c'est un modèle clair pour passer de l'idée abstraite des Neural Quantum States à une expérience reproductible, avec des métriques, des benchmarks et des observables concrets. Et pour ceux qui travaillent à l'interface des deux domaines, le guide offre une base à étendre : passer à des réseaux plus grands, ajouter des symétries, étudier l'intrication ou construire des simulations temporelles plus complexes.

Ce que cela signifie : l'approche Transformer devient peu à peu un outil de travail, pas seulement pour les tâches classiques d'IA, mais aussi pour modéliser des systèmes quantiques, où le coût d'une erreur est élevé et où les méthodes exactes s'épuisent vite. Si NetKet et JAX figurent déjà dans votre stack, ce guide fournit un point de départ quasi prêt pour des expériences de niveau recherche.