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

Technische Samenvatting: Mechanistische Interpreteerbaarheid en Causale Feature-sturing van Neurale Kwantumtoestanden via Sparse Autoencoders

Probleemstelling
Neurale Kwantumtoestanden (NQS) zijn opgekomen als krachtige variationele ansätze voor het representeren van veel-deeltjes golffuncties in de kwantummechanica, waarbij ze succes hebben geboekt in zowel de grondtoestand-representatie als de real-time dynamica van spin- en fermionische systemen. Ondanks hun empirische effectiviteit blijven NQS echter grotendeels "black boxes". Er bestaat een fundamentele kloof in het begrip van hun interne mechanismen: aangezien ze uitsluitend getraind zijn op variationele doelstellingen (zoals energieminimimalisatie) zonder expliciete optimalisatie voor specifieke fysieke observabelen, is het onduidelijk hoe NQS hoeveelheden zoals ordeparameters of correlatiefuncties nauwkeurig vastleggen. Deze opaakheid belemmert het vertrouwen in en de robuustheid van het model, met name in regimes waar numerieke benchmarks moeilijk te verkrijgen zijn. Bestaande benaderingen hebben zich grotendeels gericht op externe architecturale beperkingen of parameterelementen, waardoor er een gebrek is aan post hoc analyse met betrekking tot hoe netwerken fysieke informatie intern organiseren.

Methodologie
De auteurs stellen een systematisch kader voor met behulp van Sparse Autoencoders (SAEs), een instrument uit de mechanistische interpreteerbaarheid, om de interne activaties van NQS te deconstrueren naar sparse, interpreteerbare features.

1. NQS Architectuur: De studie maakt gebruik van autoregressieve transformer-gebaseerde NQS. De golffunctie-amplitude $\psi(\sigma)$ wordt gefactoriseerd in een Born-waarschijnlijkheid $p(\sigma)$ en een fase $\phi(\sigma)$. De transformer fungeert als een causaal sequentiemodel dat spinconfiguraties site-voor-site voorspelt.
2. SAE Toepassing: De SAE wordt toegepast op de residual stream van de laatste transformer-laag. Het leert een overcompleet woordenboek van feature-richtingen $\{f_k\}$ en sparse coëfficiënten $z_k$, zodanig dat activatievectoren $h$ kunnen worden benaderd als een sparse lineaire combinatie: $h \approx \sum z_k f_k$.
3. Ongesuperviseerde Training: Cruciaal is dat de SAE volledig ongesuperviseerd wordt getraind op de activaties van de residual stream. De SAE ontvangt geen fysieke labels (bijv. magnetisatiewaarden) en wordt uitsluitend geoptimaliseerd om activaties te reconstrueren met een sparsity-straf ($L_1$-norm).
4. Feature-Observable Correlatie: Na de training gebruiken de auteurs "mean-pooling" over de spin-keten om scalaire feature-sterktes te verkrijgen. Deze worden gecorreleerd met fysieke observabelen (bijv. magnetisatie, gestaffelde magnetisatie, half-chain correlatoren) die berekend zijn uit dezelfde gesamplede configuraties.
5. Causale Interventie (Steering): Om causaliteit vast te stellen, voeren de auteurs gerichte interventies uit. Ze schalen de activatiestuctuur van een specifieke feature ($z_{k^*} \to \alpha z_{k^*}$) terwijl anderen constant blijven, reconstrueren de gewijzigde residual stream, en observeren de resulterende verandering in de fysieke observabele. Dit gebeurt terwijl de variationele energie wordt gemonitord om te garanderen dat de interventie gelokaliseerd blijft.

Belangrijkste Bijdragen en Resultaten

• Ontdekking van Interpreteerbare Features: In de 1D Transversale-Veld Ising-model (TFIM) ontdekte de SAE autonoom sparse features die bijna perfect correleren ($|r| > 0.97$) met fysieke observabelen, waaronder magnetisatie ($M_z$), absolute magnetisatie ($|M_z|$), en half-chain correlatoren ($C_{zz}$). Deze ontdekking vond plaats over de paramagnetische fase, de ferromagnetische fase en het kritieke punt, ondanks dat de SAE geen voorafgaande kennis had van deze hoeveelheden.
• Causale Controle: De auteurs toonden aan dat deze features een causale rol spelen in de NQS-voorspellingen. Het schalen van een enkele feature-sterkte ($\alpha$) stuurde de corresponderende observabele (bijv. magnetisatie) glad en monotoon, terwijl de variationele energie nagenoeg onveranderd bleef (relatieve afwijking $< 0.02\%$). Dit bevestigt dat de features niet louter gecorreleerde artefacten zijn, maar interne coördinaten die fysieke voorspellingen aansturen.
• Real-Time Dynamica: Het kader werd uitgebreid naar real-time dynamica onder de TFIM. Een enkele statische feature-richting werd gevonden die de tijdafhankelijke magnetisatie $M_z(t)$ door de gehele evolutie-trajectie volgt. Bovendien veranderde het schalen van de activatiestuctuur van deze feature over de tijd systematisch de magnetisatie-dynamica, wat bewijst dat NQS dynamische informatie coderen in een statisch intern coördinatensysteem.
• Generalisatie naar 2D Systemen: De aanpak werd gevalideerd op het 2D Heisenberg Antiferromagnetische model (een niet-integreerbaar, 2D systeem). De SAE identificeerde een feature die sterk correleert met de gestaffelde magnetisatie (een ruimtelijk gestructureerde ordeparameter die subroosters onderscheidt). Het sturen van deze feature controleerde succesvol de gestaffelde magnetisatie, wat de toepasbaarheid van de methode buiten 1D en exact oplosbare systemen aantoont.

Significantie en Claims
Het artikel claimt dat NQS niet louter functionele benaderders van golffunctie-amplituden zijn, maar rijke, interpreteerbare interne representaties van fysieke informatie coderen.

• Diagnostisch Instrument: De aanpak biedt een methode om te beoordelen of een getrainde NQS een fysiek zinvolle representatie heeft geleerd. De aanwezigheid van features die causale controle uitoefenen op ordeparameters dient als bewijs van succesvol leren, terwijl de afwezigheid ervan een model-falen kan diagnosticeren.
• Interventie-instrument: Het werk vestigt een route voor gecontroleerde, op feature-niveau manipulatie van NQS. Door één enkele scalaire coëfficiënt in de residual stream aan te passen, kan men specifieke fysieke observabelen afstemmen zonder het gehele netwerk opnieuw te trainen of te finetunen.
• Tweerichtingsuitwisseling: De auteurs postuleren een synergetische relatie tussen kwantumveel-deeltjessystemen en mechanistische interpreteerbaarheid. NQS bieden een setting met grondwaarheid fysieke observabelen om claims over interpreteerbaarheid te testen, terwijl interpreteerbaarheidsinstrumenten methoden bieden om de "black box" van variationele golffuncties te openen, wat potentieel leidt tot betrouwbaardere en transparantere kwantumsimulaties.

De auteurs blijven bescheiden over de reikwijdte en merken op dat dit werk zich uitsluitend richt op transformer-architecturen en de residual streams van de laatste lagen, wat open vragen laat over andere architecturen (bijv. RBM's, RNN's) en laag-specifieke self-attention mechanismen.