One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective

Technische Samenvatting: One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective

Probleemstelling

Neural Quantum States (NQS) bieden een schaalbaar kader voor het benaderen van kwantum veel-deeltjes golffuncties, met name wanneer autoregressieve modellen worden gebruikt die exacte, onafhankelijke bemonstering uit de Born-distributie mogelijk maken. De optimalisatie van deze modellen blijft echter een significante flessenhals. Bestaande methoden kampen met een scherpe afruil:

• Eerstelijns optimalisatoren (bijv. Adam): Computationeel efficiënt en schaalbaar, maar negeren de geometrie van de variationele golffunctie, wat vaak leidt tot instabiele of onnauwkeurige convergentie in NQS-toepassingen.
• Stochastic Reconfiguration (SR) en varianten (bijv. minSR): Geometrisch gefundeerd (benadering van natuurlijke gradiëntafdaling), maar computationeel onhaalbaar voor grote netwerken vanwege de noodzaak om grote, vaak slecht geconditioneerde lineaire systemen op te lossen. Ze lijden ook onder numerieke fragiliteit in complexe regimes.

Hoewel autoregressieve NQS bemonsteringsproblemen oplossen, blijft stabiele en schaalbare optimalisatie voor grote autoregressieve golffunctiemodellen een open uitdaging.

Methodologie

De auteurs stellen Proximal Wavefunction Optimization (PWO) voor, een trust-region algoritme geïnspireerd door Proximal Policy Optimization (PPO) uit Reinforcement Learning (RL). De kernmethodologie rust op een formele equivalentie tussen variationele energie-minimalisatie in NQS en policy-gradient optimalisatie in RL.

1. RL-VMC Equivalentie: Het artikel stelt vast dat voor stokastische Hamiltonians de gradiënt van de variationele energie kan worden uitgedrukt als een advantage-weighted policy-gradient update over de Born-distributie. In deze mapping:

   • Spinconfiguraties komen overeen met acties.
     {% action %}
   • Gecentreerde lokale energieën ($E_{loc}(s) - \mathbb{E}[E_{loc}]$) komen overeen met advantages.
   • De Born-distributie komt overeen met de policy.

2. Dual-Channel Optimalisatie: In tegenstelling tot standaard PPO, dat één enkele waarschijnlijkheidsdistributie behandelt, zijn NQS golffuncties complex-waardig, wat een aparte behandeling van amplitude en fase vereist. PWO introduceert twee geklipte surrogate objectieven:

   • Amplitude Channel: Gebruikt een geklipte waarschijnlijkheidsratio ($r_\theta = P_\theta / P_{\theta_{old}}$) om updates te beperken, wat de standaard PPO nabootst. Dit zorgt ervoor dat de bijgewerkte golffunctie dicht bij de referentie blijft in termen van waarschijnlijkheidsmassa.
   • Phase Channel: Introduceert een nieuw geklipt surrogate dat de wrapped phase increment ($\phi_\theta$) beperkt in plaats van een waarschijnlijkheidsratio. Dit voorkomt abrupte rotaties van de golffunctie terwijl de imaginaire component van de Variational Monte Carlo (VMC) gradiënt wordt gevolgd.

3. Algoritmische Implementatie: PWO voert meerdere innerlijke updates ($K$) uit op een bevroren batch van samples getrokken uit de referentiedistributie $P_{\theta_{old}}$. Deze samples worden hergebruikt over meerdere updates door middel van importance weighting (met stop-gradient operatoren voor het phase channel) om discrepanties in de distributie te corrigeren. Dit vermijdt expliciete matrixinversie, waardoor de schaalbaarheid van eerstelijns methoden behouden blijft terwijl trust-region beperkingen worden afgedwongen.

Belangrijkste Bijdragen

• Formele Verbinding: De auteurs leggen een formele link tussen variationele energie-minimalisatie en advantage policy-gradient RL, waarbij zij aantonen dat de NQS-gradiënt een advantage-weighted vorm aanneemt over de Born-distributie.
• PWO Algoritme: Introductie van Proximal Wavefunction Optimization, een trust-region algoritme dat de waarschijnlijkheidsratio in het amplitude channel clipt en de fase-incrementen in het phase channel.
• Theoretische Garanties: Bewijs dat de PWO-surrogate voldoet aan een trust-region verbeteringsgrens en dat de eerstelijns update overeenkomt met de exacte VMC-gradiëntrichting wanneer deze on-policy is.
• Demonstratie van Schaalbaarheid: Succesvolle fine-tuning van een 1,5 miljard parameters tellende RWKV-7 Large Language Model (LLM) als een autoregressieve NQS, waarmee optimalisatie op een schaal wordt gedemonstreerd die drie ordes van grootte boven het voorgaande werk ligt.

Experimentele Resultaten

De auteurs evalueerden PWO tegenover Adam, minSR en SPRING over verschillende benchmarks:

• 1D Spin Systemen (Ising en J1–J2):
  • Ising Model: PWO convergeerde naar een relatieve fout van $10^{-7}$ in ongeveer 5 minuten, aanzienlijk sneller dan minSR (~30 minuten) en Adam.
  • Frustrated J1–J2 Chain: PWO behield een stabiele convergentie met een afnemende V-score. In contrast hiermee was minSR numeriek instabiel (6 van de 10 runs produceerden NaNs), en Adam plateauerde op een veel hogere fout.
• 2D Frustrated J1–J2 Model: Op een $10 \times 10$ rooster bereikte PWO lagere energieniveaus sneller dan Adam en behield een lagere variantie-gebaseerde foutsignaal (V-score) binnen hetzelfde wall-clock budget.
• Schaalbaarheid: Naarmate de modelgrootte toenam (van "tiny" naar "medium" NQS), demonstreerde PWO consistente stabiliteit en monotone energieverbetering. Adam degradeerde in stabiliteit en efficiëntie voor grotere modellen, terwijl minSR niet in staat was competitieve energiewaarden te bereiken.
• Grootschalige Fine-Tuning: Het 1,5B-parameter RWKV-7 model, gefinetuned met PWO, bereikte lagere uiteindelijke relatieve fouten en V-scores dan Adam, waarbij een stabiele training werd behouden gedurende het gehele budget.

Betekenis en Claims

Het artikel beweert dat PWO de kloof overbrugt tussen de schaalbaarheid van eerstelijns optimalisatoren en de stabiliteit van geometrische methoden (zoals SR) door gebruik te maken van de trust-region principes van RL.

• Stabiliteit in Moeilijke Regimes: De auteurs stellen dat PPO-achtige proximale objectieven de noodzakelijke stabiliteit bieden om expressieve NQS te trainen op moeilijke, gefrustreerde kwantum Hamiltonians waar traditionele methoden falen of instabiel worden.
• Sample Efficiëntie: Door het mogelijk te maken om samples herhaaldelijk te gebruiken voor meerdere updates zonder expliciete matrixinversie, verbetert PWO de wall-clock convergentiesnelheid.
• Schaal: De succesvolle toepassing op een 1,5B-parameter model suggereert dat NQS-optimalisatie nu het regime van moderne grootschalige sequentiemodellering kan betreden. De auteurs stellen dat dit de deur opent naar onderzoek naar de vraag of schaal alleen de huidige barrières in kwantum veel-deeltjes simulatie kan overwinnen, in lijn met de "bitter lesson" in AI.

Het werk beweert niet dat LLM's de optimale inductieve bias zijn voor 1D Hamiltonians, maar gebruikt het RWKV-7 model juist om de effectiviteit van het proximale objectief op ongekende schaal aan te tonen.