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

Technische Zusammenfassung: Mechanistische Interpretierbarkeit und kausale Steuerung von Merkmalen Quantenzuständen mittels Sparse Autoencoders

Problemstellung
Neuronale Quantenzustände (NQS) haben sich als leistungsfähige variationale Ansätze zur Repräsentation von Quanten-Vielteilchen-Wellenfunktionen etabliert und erzielen Erfolge sowohl bei der Darstellung von Grundzuständen als auch in der Echtzeitdynamik bei Spin- und Fermionsystemen. Trotz ihrer empirischen Effizienz bleiben NQS jedoch weitgehend „Black Boxes“. Es besteht eine fundamentale Lücke im Verständnis ihrer internen Mechanismen: Da sie ausschließlich auf variationalen Zielen (wie etwa der Energieminimierung) trainiert werden, ohne explizit für spezifische physikalische Observablen optimiert zu sein, ist unklar, wie NQS Größen wie Ordnungsparameter oder Korrelationsfunktionen präzise erfassen. Diese Undurchsichtigkeit behindert das Vertrauen in das Modell und dessen Robustheit, insbesondere in Regimen, in denen numerische Benchmarks schwer zu erhalten sind. Bestehende Ansätze konzentrierten sich weitgehend auf externe architektonische Beschränkungen oder Parameteranalysen, wodurch eine Analyse post hoc darüber fehlt, wie Netzwerke interne physikalische Informationen organisieren.

Methodik
Die Autoren schlagen ein systematisches Framework vor, das Sparse Autoencoder (SAEs) verwendet – ein Werkzeug aus der mechanistischen Interpretierbarkeit –, um die internen Aktivierungen von NQS in dünnbesetzte (sparse), interpretierbare Merkmale zu zerlegen.

1. NQS-Architektur: Die Studie nutzt autoregressive, Transformer-basierte NQS. Die Wellenfunktionsamplitude $\psi(\sigma)$ wird in eine Born-Wahrscheinlichkeit $p(\sigma)$ und eine Phase $\phi(\sigma)$ faktorisiert. Der Transformer fungiert als kausales Sequenzmodell, das Spin-Konfigurationen Schritt für Schritt vorhersagt.
2. SAE-Anwendung: Der SAE wird auf den Residual Stream (Residuenstrom) der finalen Transformer-Schicht angewendet. Er lernt ein überkomplettes Dictionary von Merkmalsrichtungen $\{f_k\}$ und dünnbesetzten Koeffizienten $z_k$, sodass der Aktivierungsvektor $h$ als dünnbesetzte Linearkombination approximiert werden kann: $h \approx \sum z_k f_k$.
3. Unüberwachtes Training: Entscheidend ist, dass der SAE vollständig unüberwacht auf den Aktivierungen des Residual Streams trainiert wird. Er erhält keine physikalischen Labels (z. B. Magnetisierungswerte) und wird ausschließlich darauf optimiert, die Aktivierungen unter Anwendung eines Sparsitäts-Penalties ($L_1$-Norm) zu rekonstruieren.
4. Merkmals-Observable-Korrelation: Nach dem Training führen die Autoren ein „Mean-Pooling“ der dünnbesetzten Koeffizienten über die Spin-Kette durch, um skalare Merkmalsstärken zu erhalten. Diese werden mit physikalischen Observablen (z. B. Magnetisierung, gestaffelte Magnetisierung, Half-Chain-Korrelatoren) korreliert, die aus denselben Stichproben-Konfigurationen berechnet wurden.
5. Kausale Intervention (Steuerung): Um Kausalität zu etablieren, führen die Autoren gezielte Interventionen durch. Sie skalieren die Aktivierungsstärke eines spezifischen Merkmals ($z_{k^*} \to \alpha z_{k^*}$), während andere fixiert bleiben, rekonstruieren den modifizierten Residual Stream und beobachten die resultierende Änderung der physikalischen Observable. Dies geschieht unter Überwachung der variationalen Energie, um sicherzustellen, dass die Intervention lokalisiert bleibt.

Wesentliche Beiträge und Ergebnisse

• Entdeckung interpretierbarer Merkmale: Im 1D-Transversalfeld-Ising-Modell (TFIM) entdeckte der SAE autonom dünnbesetzte Merkmale, die nahezu perfekt mit physikalischen Observablen korrelieren ($|r| > 0,97$), einschließlich Magnetisierung ($M_z$), absoluter Magnetisierung ($|M_z|$) und Half-Chain-Korrelatoren ($C_{zz}$). Diese Entdeckung erfolgte über die paramagnetische Phase, die ferromagnetische Phase und den kritischen Punkt hinweg, obwohl der SAE keinerlei Vorwissen über diese Größen besaß.
• Kausale Kontrolle: Die Autoren zeigten, dass diese Merkmale eine kausale Rolle bei den NQS-Vorhersagen spielen. Das Skalieren einer einzelnen Merkmalsstärke ($\alpha$) steuerte die entsprechende Observable (z. B. Magnetisierung) glatt und monoton, während die variationale Energie nahezu unverändert blieb (relative Abweichung $< 0,02 \%$). Dies bestätigt, dass die Merkmale nicht bloß korrelierte Artefakte sind, sondern interne Koordinaten, welche die physikalischen Vorhersagen kontrollieren.
• Echtzeitdynamik: Das Framework wurde auf die Echtzeitdynamik unter dem TFIM ausgeweitet. Eine einzige statische Merkmalsrichtung konnte die zeitabhängige Magnetisierung $M_z(t)$ über die gesamte Evolutions-Trajektorie hinweg verfolgen. Darüber hinaus änderte das Skalieren der Merkmalsaktivierung über die Zeit systematisch die Magnetisierungsdynamik, was beweist, dass NQS dynamische Informationen in einem statischen internen Koordinatensystem kodieren.
• Generalisierung auf 2D-Systeme: Der Ansatz wurde am 2D-Heisenberg-Antiferromagnet (ein nicht-integrierbares 2D-System) validiert. Der SAE identifizierte ein Merkmal, das stark mit der gestaffelten Magnetisierung (einem räumlich strukturierten Ordnungsparameter, der Subgitter unterscheidet) korreliert. Die Steuerung dieses Merkmals kontrollierte erfolgreich die gestaffelte Magnetisierung, was die Anwendbarkeit der Methode über 1D- und exakt lösbare Systeme hinaus demonstriert.

Bedeutung und Behauptungen
Das Paper behauptet, dass NQS nicht lediglich funktionale Approximatoren von Wellenfunktionsamplituden sind, sondern reiche, interpretierbare interne Repräsentationen physikalischer Informationen kodieren.

• Diagnostisches Werkzeug: Der Ansatz bietet eine Methode, um zu beurteilen, ob ein trainiertes NQS eine physikalisch sinnvolle Repräsentation erlernt hat. Das Vorhandensein von Merkmalen, die Ordnungsparameter kausal kontrollieren, dient als Beleg für erfolgreiches Lernen, während deren Fehlen ein Modellversagen diagnostizieren könnte.
• Interventionswerkzeug: Die Arbeit etabliert einen Weg für die kontrollierte Manipulation auf Merkmalsebene von NQS. Durch Anpassung eines einzelnen skalaren Koeffizienten im Residual Stream kann man spezifische physikalische Observablen tunen, ohne das gesamte Netzwerk neu zu trainieren oder feinabzustimmen.
• Zweiseitiger Austausch: Die Autoren postulieren eine synergistische Beziehung zwischen der Quanten-Vielteilchenphysik und der mechanistischen Interpretierbarkeit. NQS bieten ein Setting mit physikalischen Grundwahrheiten (Ground-Truth), um Interpretierbarkeitsansprüche zu testen, während Interpretierbarkeitstools Methoden bieten, um die „Black Box“ variationaler Wellenfunktionen zu öffnen, was potenziell zu zuverlässigeren und transparenteren Quantensimulationen führt.

Die Autoren bleiben hinsichtlich des Umfangs bescheiden und merken an, dass sich diese Arbeit ausschließlich auf Transformer-Architekturen und die finalen Layer-Residual-Streams konzentriert, was Fragen bezüglich anderer Architekturen (z. B. RBMs, RNNs) und layer-gelöster Self-Attention-Mechanismen offen lässt.