Mechanistic Interpretability and Causal Feature Steering of Neural Quantum States via Sparse Autoencoders

Résumé Technique : Interprétabilité Mécaniste et Pilotage de Caractéristiques Causales d'États Quantiques Neuronaux via des Autoencodeurs Creux

Énoncé du Problème
Les États Quantiques Neuronaux (NQS) sont apparus comme de puissants ansatz variationnels pour représenter les fonctions d'onde de systèmes à plusieurs corps quantiques, rencontrant un succès tant dans la représentation de l'état fondamental que dans la dynamique en temps réel pour les systèmes de spins et de fermions. Cependant, malgré leur efficacité empirique, les NQS restent largement des « boîtes noires ». Un fossé fondamental subsiste dans la compréhension de leurs mécanismes internes : étant entraînés uniquement sur des objectifs variationnels (tels que la minimisation de l'énergie) sans optimisation explicite pour des observables physiques spécifiques, il est difficile de savoir comment les NQS capturent avec précision des quantités telles que les paramètres d'ordre ou les fonctions de corrélation. Cette opacité entrave la confiance et la robustesse des modèles, particulièrement dans les régimes où les références numériques sont difficiles à obtenir. Les approches existantes se sont largement concentrées sur les contraintes architecturales externes ou l'analyse de paramètres, laissant un manque d'analyse a posteriori concernant la manière dont les réseaux organisent l'information physique en interne.

Méthodologie
Les auteurs proposent un cadre systématique utilisant des Autoencodeurs Creux (SAE), un outil issu de l'interprétabilité mécaniste, pour décomposer les activations internes des NQS en caractéristiques creuses et interprétables.

1. Architecture NQS : L'étude utilise des NQS de type transformeur autorégressif. L'amplitude de la fonction d'onde $\psi(\sigma)$ est factorisée en une probabilité de Born $p(\sigma)$ et une phase $\phi(\sigma)$. Le transformeur agit comme un modèle de séquence causal, prédisant les configurations de spins site par site.
2. Application du SAE : Le SAE est appliqué au flux résiduel (residual stream) de la dernière couche du transformeur. Il apprend un dictionnaire surcomplet de directions de caractéristiques $\{f_k\}$ et de coefficients creux $z_k$ tels que les vecteurs d'activation $h$ sont approximés par une combinaison linéaire creuse : $h \approx \sum z_k f_k$.
3. Entraînement Non Supervisé : Crucialement, le SAE est entraîné entièrement de manière non supervisée sur les activations du flux résiduel. Il ne reçoit aucune étiquette physique (ex: valeurs de magnétisation) et est optimisé uniquement pour reconstruire les activations avec une pénalité de parcimonie (norme $L_1$).
4. Corrélation Caractéristique-Observable : Après l'entraînement, les auteurs effectuent un « regroupement moyen » (mean-pooling) des coefficients creux sur la chaîne de spins pour obtenir des intensités de caractéristiques scalaires. Ceux-ci sont corrélés avec des observables physiques (ex: magnétisation, magnétisation décalée, corrélateurs de demi-chaîne) calculés à partir des mêmes configurations échantillonnées.
5. Intervention Causale (Pilotage) : Pour établir la causalité, les auteurs effectuent des interventions ciblées. Ils redimensionnent l'intensité d'activation d'une caractéristique spécifique ($z_{k^*} \to \alpha z_{k^*}$) tout en maintenant les autres fixes, reconstruisent le flux résiduel modifié, et observent le changement résultant de l'observable physique. Cela est effectué en surveillant l'énergie variationnelle pour s'assurer que l'intervention reste localisée.

Contributions Clés et Résultats

• Découverte de Caractéristiques Interprétables : Dans le modèle d'Ising en champ transverse (TFIM) en 1D, le SAE a découvert de manière autonome des caractéristiques creuses qui sont presque parfaitement corrélées ($|r| > 0,97$) avec des observables physiques, incluant la magnétisation ($M_z$), la magnétisation absolue ($|M_z|$), et les corrélateurs de demi-chaîne ($C_{zz}$). Cette découverte a eu lieu à travers la phase paramagnétique, la phase ferromagnétique et le point critique, bien que le SAE n'ait aucune connaissance préalable de ces quantités.
• Contrôle Causal : Les auteurs ont démontré que ces caractéristiques jouent un rôle causal dans les prédictions des NQS. Le redimensionnement de l'intensité d'une caractéristique ($\alpha$) pilote de manière fluide et monotone l'observable correspondante (ex: la magnétisation) tout en laissant l'énergie variationnelle pratiquement inchangée (écart relatif $< 0,02 \%$). Cela confirme que les caractéristiques ne sont pas de simples artefacts de corrélation, mais des coordonnées internes qui contrôlent les prédictions physiques.
• Dynamique en Temps Réel : Le cadre a été étendu à la dynamique en temps réel sous le TFIM. Une direction de caractéristique statique unique a été capable de suivre la magnétisation dépendante du temps $M_z(t)$ tout au long de la trajectoire d'évolution complète. De plus, le redimensionnement de l'intensité d'activation de cette caractéristique à travers le temps a systématiquement altéré la dynamique de la magnétisation, prouvant que les NQS encodent l'information dynamique dans un système de coordonnées internes statiques.
• Généralisation aux Systèmes 2D : L'approche a été validée sur le modèle de Heisenberg antiferromagnétique en 2D (un système non intégrable en 2D). Le SAE a identifié une caractéristique fortement corrélée avec la magnétisation décalée (un paramètre d'ordre spatialement structuré distinguant les sous-réseaux). Le pilotage de cette caractéristique a permis de contrôler avec succès la magnétisation décalée, démontant l'applicabilité de la méthode au-delà des systèmes 1D et de solutions exactes.

Signification et Revendications
L'article affirme que les NQS ne sont pas seulement des approximateurs fonctionnels d'amplitudes de fonctions d'onde, mais qu'ils encodent des représentations internes riches et interprétables de l'information physique.

• Outil de Diagnostic : Cette approche fournit une méthode pour évaluer si un NQS entraîné a appris une représentation physiquement cohérente. La présence de caractéristiques qui contrôlent causalement les paramètres d'ordre sert de preuve d'un apprentissage réussi, tandis que leur absence pourrait diagnostiquer un échec du modèle.
• Outil d'Intervention : Ce travail établit une voie pour la manipulation contrôlée au niveau des caractéristiques des NQS. En ajustant un seul coefficient scalaire dans le flux résiduel, on peut régler des observables physiques spécifiques sans réentraîner ou affiner l'ensemble du réseau.
• Échange Bidirectionnel : Les auteurs postulent une relation synergique entre la physique des systèmes à plusieurs corps quantiques et l'interprétabilité mécaniste. Les NQS fournissent un cadre avec des observables physiques de référence pour tester les affirmations d'interprétabilité, tandis que les outils d'interprétabilité offrent des méthodes pour ouvrir la « boîte noire » des fonctions d'onde variationnelles, menant potentiellement à des simulations quantiques plus fiables et transparentes.

Les auteurs restent modestes quant à la portée de l'étude, notant que ce travail se concentre exclusivement sur les architectures de type transformeur et les flux résiduels de la dernière couche, laissant ouvertes des questions concernant d'autres architectures (ex: RBM, RNN) et les mécanismes d'auto-attention résolus par couche.