Parallel Scan Recurrent Neural Quantum States: Making RNNs Scalable for Quantum Many-Body Problems

Paper: 2605.13807 Authors: Ejaaz Merali, Mohamed Hibat-Allah, Mohammad Kohandel, Richard T. Scalettar, Ehsan Khatami Categories: cond-mat.str-el, cond-mat.dis-nn, cs.LG, physics.comp-ph, quant-ph

The Gap

Neural quantum states (NQS) use neural networks to represent quantum wave functions for many-body systems. Transformers dominate because they parallelize well across GPUs. Recurrent neural networks (RNNs) were dismissed as inherently sequential—you need the hidden state from position i to compute position i+1, so no parallelism, right? This created a perceived architectural ceiling: if you want to scale to large lattices (say, 50×50 spins), you need transformers with their quadratic memory cost and massive compute budgets.

The gap: RNNs have linear memory scaling and are conceptually simpler, but everyone assumed the sequential bottleneck made them impractical for large quantum systems. No one seriously tried to make RNNs competitive at scale.

Problem: Quantum many-body ground states on large lattices
   |
   v
Assumption: RNNs = sequential = slow = can't scale
   |
   v
Method: Apply parallel scan algorithms to RNN recurrence
   |
   v
Evidence: 52x52 lattices, matches QMC benchmarks, modest GPU hours
   |
   v
Conclusion: RNNs are viable for large-scale NQS with less compute

The Increment

One sentence: Before this paper, scaling neural quantum states meant using transformers with quadratic memory; after, RNNs with parallel scans achieve comparable accuracy on 2704-spin systems with linear memory and accessible compute.

Core Mechanism

The method has three layers. First, an autoregressive RNN wave function that generates spin configurations one site at a time, maintaining a hidden state that encodes the quantum correlations seen so far. Second, a parallel scan algorithm (specifically, they use associative scans over the recurrence relation) that computes all hidden states simultaneously instead of sequentially. Third, variational Monte Carlo training where you sample configurations from the RNN, compute local energies, and update parameters to minimize energy.

Data flows like this: you feed in a partial spin configuration (say, spins 1 through k). The RNN’s recurrence relation updates the hidden state at each position. Normally you’d compute h_1, then h_2, then h_3 sequentially. The parallel scan reformulates this as a tree of operations that can run in log(n) depth on parallel hardware. Once you have all hidden states, the RNN outputs probabilities for the next spin. During training, you sample full configurations, evaluate the quantum Hamiltonian, and backpropagate through the parallel scan to update weights.

Input spins:  s_1  s_2  s_3  s_4  ...  s_n
                |    |    |    |         |
                v    v    v    v         v
Parallel    [  h_1  h_2  h_3  h_4  ...  h_n  ]  <- computed in log(n) steps
Scan:          /  \  /  \  /  \            via associative tree reduction
              /    \/    \/    \
             /      \    /      \
            +--------+--+--------+
                     |
                     v
Output:         P(s_next | s_1...s_n)
                     |
                     v
Sample -> Evaluate H -> Backprop -> Update weights

Think of it like a factory assembly line that’s been redesigned. The old RNN is a single-file conveyor belt: each worker (time step) must wait for the previous worker to finish before starting. The parallel scan is like reorganizing the factory into a binary tree of workstations. Workers at the leaves start simultaneously, pass partial results up to their parents, who combine them and pass further up, until the root has the final result. The total work is the same, but the critical path (longest dependency chain) shrinks from n steps to log(n) steps. For quantum states, each “worker” is computing how quantum correlations propagate through the lattice, and the tree structure lets you compute all positions’ correlations in parallel while preserving the sequential dependencies mathematically.

Key Concepts

• Autoregressive wave function: In quantum mechanics, a wave function assigns an amplitude (complex number) to every possible configuration of the system. For n spins, that’s 2^n configurations—exponentially large. An autoregressive model factors this into a product: ψ(s_1, s_2, …, s_n) = P(s_1) × P(s_2|s_1) × P(s_3|s_1,s_2) × … Each factor is a conditional probability the RNN can compute. Instead of storing 2^n numbers, you store RNN weights. When you need ψ for a specific configuration, you run the RNN forward to compute the product. It’s like compressing a phone book by storing the algorithm that generates names instead of the names themselves.

• Parallel scan (associative scan): Imagine you want to compute cumulative sums: given [a, b, c, d], output [a, a+b, a+b+c, a+b+c+d]. Sequentially, that’s 4 steps. But addition is associative: (a+b)+c = a+(b+c). So you can compute (a+b) and (c+d) in parallel, then add those results. This generalizes to any associative operation. RNN recurrence h_t = f(h_{t-1}, x_t) can be rewritten as an associative operation if f has the right structure. The parallel scan builds a binary tree: leaves compute pairwise combinations, parents combine their children’s results, and the root has the final answer. Depth is log(n) instead of n.

• Variational Monte Carlo (VMC): You want the ground state (lowest energy configuration) of a quantum Hamiltonian H. VMC treats this as optimization: parameterize the wave function ψ(s; θ) with neural network weights θ, sample configurations s from |ψ|^2 (the probability distribution), compute the local energy E_loc(s) = (Hψ)(s)/ψ(s) for each sample, and minimize the average energy ⟨E_loc⟩ by gradient descent on θ. It’s like finding the lowest point in a landscape by randomly dropping marbles (samples) and rolling downhill (gradient descent). The RNN gives you ψ, the parallel scan makes sampling fast, and VMC ties it all together into a training loop.

Framework Shift

Before (Transformers):                After (PSR-NQS):
                                      
Spin lattice (n sites)                Spin lattice (n sites)
       |                                     |
       v                                     v
Transformer encoder                   RNN with parallel scan
  - Attention: O(n^2) memory            - Recurrence: O(n) memory
  - All-to-all interactions             - Sequential structure
  - Massive parallelism                 - Log-depth parallelism
       |                                     |
       v                                     v
Sample configurations                 Sample configurations
       |                                     |
       v                                     v
VMC training (GPU cluster)            VMC training (single GPU)

From architecture-driven scaling (throw more compute at transformers) to algorithm-driven scaling (make RNNs parallel via scan primitives), the core shift is exploiting mathematical structure instead of brute-force parallelism.

Expert Assessment

Problem choice: Real gap. The transformer monoculture in NQS is partly inertia—RNNs were written off without exploring modern recurrent architectures. The authors correctly identified that parallel scans, which have been around since the 1980s in parallel computing, were never seriously applied to quantum state representations. This sits at the intersection of two mature fields (parallel algorithms and NQS), which is often where low-hanging fruit hides.

Method maturity: Clever application of known techniques rather than fundamental innovation. Parallel scans are textbook material; the insight is recognizing that RNN recurrence fits the associative structure needed for scans. The iterative retraining scheme (train on small lattice, use as initialization for larger lattice) is standard transfer learning. Nothing here is conceptually hard, which is actually a strength—it’s engineering that should have been done years ago.

Experimental integrity: Baselines are fair. They compare against exact diagonalization (small systems), quantum Monte Carlo (gold standard for these models), and prior transformer-based NQS. The 52×52 result is impressive but comes with a caveat: they use iterative retraining, which means the model isn’t learning from scratch at that scale. The energy errors are small (within QMC error bars), but they don’t report variance or multiple runs, so it’s hard to assess stability. The compute budget comparison (single GPU vs. cluster) is compelling but lacks detail—what GPU, how many hours, what batch size?

Writing quality: The paper front-loads motivation well but buries implementation details. Section III (Methods) is too terse—the parallel scan description assumes familiarity with the algorithm and doesn’t explain how the RNN recurrence is reformulated to be associative. The results section would benefit from ablations: how much does iterative retraining contribute vs. the parallel scan itself? The related work undersells prior RNN-based NQS efforts, making the contribution seem more novel than it is.

Verdict: weak accept — Solid engineering contribution that challenges a false dichotomy (RNNs vs. transformers) in an important application domain, but the experimental evaluation could be more rigorous and the writing clearer.

Takeaways

If you’re working with autoregressive models on sequences where you need both training efficiency and inference speed, check if your recurrence is associative (or can be approximated as such). Parallel scans are underutilized outside of niche parallel computing circles. The iterative retraining trick—train small, transfer to large—is a practical way to scale when direct training is prohibitive. More broadly, when a field converges on one architecture (transformers, in this case), it’s worth revisiting dismissed alternatives with modern tools. The gap between “theoretically possible” and “actually implemented” is often just someone bothering to write the code.