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

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

Definizione del Problema

Gli Stati Quantistici Neurali (NQS) offrono un framework scalabile per approssimare le funzioni d'onda di molti corpi quantistici, in particolare quando si utilizzano modelli autoregressivi che consentono il campionamento esatto e indipendente dalla distribuzione di Born. Tuttavia, l'ottimizzazione di questi modelli rimane un collo di bottiglia significativo. I metodi esistenti affrontano un netto compromesso:

• Ottimizzatori del primo ordine (es. Adam): Computazionalmente efficienti e scalabili, ma ignorano la geometria della funzione d'onda variazionale, portando spesso a una convergenza instabile o imprecisa nelle applicazioni NQS.
• Stochastic Reconfiguration (SR) e varianti (es. minSR): Geometricamente fondati (approssimando il gradiente naturale), ma computazionalmente proibitivi per reti di grandi dimensioni a causa della necessità di risolvere sistemi lineari ampi e spesso mal condizionati. Soffrono inoltre di fragilità numerica in regimi complessi.

Sebbene gli NQS autoregressivi risolvano i colli di bottiglia del campionamento, l'ottimizzazione stabile e scalabile per grandi modelli di funzioni d'onda autoregressive rimane una sfida aperta.

Metodologia

Gli autori propongono la Proximal Waveframe Optimization (PWO), un algoritmo a regione di fiducia (trust-region) ispirato al Proximal Policy Optimization (PPO) del Reinforcement Learning (RL). La metodologia centrale si basa su un'equivalenza formale tra la minimizzazione dell'energia variazionale in NQS e l'ottimizzazione del policy-gradient in RL.

1. Equivalenza RL-VMC: Il paper stabilisce che per gli Hamiltoniani stoquastici, il gradiente dell'energia variazionale può essere espresso come un aggiornamento di tipo advantage-weighted policy-gradient sulla distribuzione di Born. In questa mappatura:

   • Le configurazioni di spin corrispondono ad azioni.
   • Le energie locali centrate ($E_{loc}(s) - \mathbb{E}[E_{loc}]$) corrispondono ad advantage.
   • La distribuzione di Born corrisponde alla policy.

2. Ottimizzazione a Canale Doppio: A differenza del PPO standard che gestisce una singola distribuzione di probabilità, le funzioni d'onda NQS sono complesse, richiedendo una gestione separata di ampiezza e fase. PWO introduce due obiettivi surrogati con clipping:

   • Canale di Ampiezza: Utilizza un rapporto di probabilità clippato ($r_\theta = P_\theta / P_{\theta_{old}}$) per vincolare gli aggiornamenti, rispecchiando il PPO standard. Ciò assicura che la funzione d'onda aggiornata rimanga vicina al riferimento in termini di massa di probabilità.
   • Canale di Fase: Introduce un nuovo surrogato clippato che vincola l'incremento della fase avvolta ($\phi_\theta$) piuttosto che un rapporto di probabilità. Questo evita rotazioni brusche della funzione d'onda pur seguendo la parte immaginaria del gradiente di Variational Monte Carlo (VMC).

3. Implementazione Algoritmica: PWO esegue molteplici aggiornamenti interni ($K$) su un batch congelato di campioni estratti dalla distribuzione di riferimento $P_{\theta_{old}}$. Riutilizza questi campioni attraverso vari aggiornamenti impiegando il pesatura per importanza (con operatori stop-gradient per il canale di fase) per correggere i disallineamenti della distribuzione. Ciò evita l'inversione esplicita di matrici, mantenendo la scalabilità dei metodi del primo ordine pur imponendo vincoli di regione di fiducia.

Contributi Chiave

• Connessione Formale: Gli autori collegano formalmente la minimizzazione dell'energia variazionale all'ottimizzazione del policy-gradient di tipo advantage, dimostrando che il gradiente NQS ammette una forma advantage-weighted sulla distribuzione di Born.
• Algoritmo PWO: Introduzione della Proximal Wavefunction Optimization, un algoritmo a regione di fiducia che applica il clipping al rapporto di probabilità nel canale di ampiezza e agli incrementi di fase nel canale di fase.
• Garanzie Teoriche: Dimostrazione che il surrogato PWO soddisfa un limite di miglioramento della regione di fiducia e che l'aggiornamento del primo ordine corrisponde alla direzione esatta del gradiente VMC quando è on-policy.
• Dimostrazione di Scalabilità: Fine-tuning riuscito di un modello Large Language Model (LLM) RWKV-7 da 1,5 miliardi di parametri come NQS autoregressivo, dimostrando l'ottimizzazione a una scala superiore di tre ordini di grandezza rispetto ai lavori precedenti.

Risultati Sperimentali

Gli autori hanno valutato PWO contro Adam, minSR e SPRING attraverso diversi benchmark:

• Sistemi di Spin 1D (Ising e J1–J2):
  • Modello Ising: PWO è convergito a un errore relativo di $10^{-7}$ in circa 5 minuti, significativamente più velocemente di minSR (~30 minuti) e Adam.
  • Catena frustrata J1–J2: PWO ha mantenuto una convergenza stabile con un V-score decrescente. Al contrario, minSR è stato numericamente instabile (6 run su 10 hanno prodotto NaN) e Adam si è stabilizzato su un errore molto più elevato.
• Modello 2D Frustrato J1–J2: Su un reticolo $10 \times 10$, PWO ha raggiunto regimi di energia inferiore più velocemente di Adam e ha mantenuto un segnale di errore basato sulla varianza (V-score) più basso entro lo stesso budget di tempo di calcolo (wall-clock).
• Scalabilità: All'aumentare della dimensione del modello (da "tiny" a "medium" NQS), PWO ha dimostrato una stabilità costante e un miglioramento energetico monotono. Adam ha degradato in stabilità ed efficienza per modelli più grandi, mentre minSR non è riuscito a raggiungere valori di energia competitivi.
• Fine-Tuning su Larga Scala: Il modello RWKV-7 da 1,5 miliardi di parametri, rifinito con PWO, ha ottenuto errori relativi e V-score finali inferiori rispetto ad Adam, mantenendo una formazione stabile durante tutto il budget.

Significato e Rivendicazioni

Il paper sostiene che PWO colmi il divario tra la scalabilità degli ottimizzatori del primo ordine e la stabilità dei metodi geometrici (come SR) sfruttando i principi di regione di fiducia del Reinforcement Learning.

• Stabilità in Regimi Difficili: Gli autori affermano che gli obiettivi prossimali in stile PPO fornisiscono la stabilità necessaria per addestrare NQS espressivi su Hamiltoniani quantistici difficili e frustrati dove i metodi tradizionali falliscono o diventano instabili.
• Efficienza di Campionamento: Consentendo il riutilizzo dei campioni attraverso molteplici aggiornamenti senza l'inversione esplicita di matrici, PWO migliora la velocità di convergenza in termini di tempo di calcolo (wall-clock).
• Scala: Il successo dell'applicazione a un modello da 1,5 miliardi di parametri suggerisce che l'ottimizzazione degli NQS può ora entrare nel regime della moderna modellazione di sequenze su larga scala. Gli autori postulano che ciò apra la strada all'indagine su come la sola scala possa superare le attuali barriere nella simulazione del molti corpo quantistico, allineandosi alla "lezione amara" (bitter lesson) dell'IA.

Il lavoro non pretende che gli LLM siano il bias induttivo ottimale per gli Hamiltoniani 1D, ma utilizza il modello RWKV-7 per dimostrare l'efficacia dell'obiettivo prossimale a scale senza precedenti.