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

Les architectures Transformer commencent à trouver leur place en physique quantique computationnelle : un nouveau guide pratique montre comment assembler un pipeline complet de Neural Quantum States en utilisant NetKet et JAX pour un problème complexe — la chaîne de spins J1-J2 de Heisenberg frustrée. Ce n'est pas une théorie abstraite, mais un cadre reproductible où le modèle, l'échantillonneur, l'optimisation et la vérification de la précision sont réunis dans une seule boucle de recherche. L'idée centrale du guide est que l'architecture Transformer est bien adaptée à la description des états quantiques à plusieurs corps, où les corrélations à longue portée entre les particules sont importantes.

Les méthodes numériques conventionnelles se heurtent rapidement à la malédiction de la dimensionnalité de l'espace des états, surtout si le système est frustré, c'est-à-dire que les interactions concurrentes l'empêchent d'atteindre un minimum d'énergie ordonné simple. Dans ces conditions, Neural Quantum States permettent que la fonction d'onde soit représentée sous la forme d'un réseau de neurones paramétrés, qui est ensuite optimisé par Monte Carlo Variationnel. NetKet agit comme un environnement prêt pour les calculs quantiques, tandis que JAX est le moteur pour l'optimisation accélérée de haute précision.

Le guide configure d'abord la partie physique de base du problème. 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 leurs voisins suivants avec J2. C'est précisément cette combinaison qui crée la frustration rendant le problème intéressant et non trivial.

Pour décrire le système, on utilise un graphe dans NetKet, un espace de Hilbert de particules de spin 1/2 avec projection totale fixe, et un opérateur hamiltonien assemblé via GraphOperator. En parallèle, la précision 64 bits est activée dans JAX, ce qui est essentiel pour des calculs stables dans cette classe de problèmes. Commence alors l'apprentissage automatique sous sa forme pure.

La fonction d'onde est définie par un modèle TransformerLogPsi personnalisé sur Flax : les configurations de spins sont codées sous forme de tokens, reçoivent des plongements et des représentations positionnelles, puis passent par plusieurs blocs d'auto-attention et des couches feed-forward. L'exemple utilise une dimension cachée de 96, quatre têtes d'attention et six couches Transformer. Le modèle retourne le logarithme complexe de l'amplitude de la fonction d'onde — ceci est critique car un état quantique ne peut pas être adéquatement décrit 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 ansatze locaux. Une valeur particulière du guide est 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 : un échantillonneur MetropolisExchange, un état variationnel MCState, un optimiseur Adam et Stochastic Reconfiguration comme analogue de la descente de gradient naturel pour les états quantiques.

La configuration de l'exemple utilise 4096 échantillons, rejet des états initiaux dans les chaînes et environ 250 itérations d'optimisation. Un tel stack est nécessaire non seulement pour atteindre une énergie basse, mais aussi pour contrôler la convergence. Le code enregistre les trajectoires d'énergie moyenne et de variance, pour que l'on puisse voir à quel point le modèle se meut stabilement vers une bonne solution.

Après l'entraînement, le pipeline est utilisé comme outil de recherche. L'auteur effectue des calculs pour plusieurs valeurs de J2 dans la plage de 0 à 0,7 pour une chaîne de longueur 24 nœuds, enregistre les énergies finales et estime le pic du facteur de structure. Cela permet non seulement d'ajuster les paramètres du réseau de neurones, mais d'observer comment le comportement physique du système change lorsque la frustration augmente et où les transitions entre différentes phases d'ordre magnétique peuvent apparaître.

Pour une vérification supplémentaire de la qualité, le modèle est comparé à la diagonalisation exacte sur un système plus petit de taille 14 nœuds à l'aide de la méthode de Lanczos. La comparaison des énergies fournit une référence numérique claire : à quel point le Transformer variationnel est proche de la solution exacte où le calcul exact est encore possible. Le sens pratique du matériel est qu'il comble le fossé entre deux mondes — les architectures modernes de l'apprentissage profond et les problèmes réels de la physique computationnelle.

Pour les ingénieurs ML, c'est un bon exemple de la façon dont un Transformer peut être utilisé en dehors du texte, des images et des données tabulaires standard. Pour les physiciens, c'est un modèle compréhensible pour passer de l'idée abstraite de Neural Quantum States à une expérience reproductible avec des métriques concrètes, des benchmarks et des quantités observables. Et pour ceux qui travaillent à l'intersection de ces domaines, le guide fournit une base qui peut être étendue : 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 graduellement un outil de travail non seulement pour les tâches classiques d'IA, mais aussi pour la modélisation de systèmes quantiques, où le coût de l'erreur est élevé et les méthodes exactes s'épuisent rapidement. Si NetKet et JAX sont déjà dans votre stack de travail, ce matériel fournit un point de départ pratiquement prêt pour des expériences de niveau recherche.