NQS-Agent: Health-Aware Agentic Hyperparameter Optimization for Neural-Network Quantum States

Abstract: Neural-network quantum states (NQS) provide expressive variational representations for strongly correlated quantum many-body systems, but their practical accuracy depends sensitively on architecture-level hyperparameters and optimization schedules. Here we develop NQS-Agent, an implemented open-source software framework for health-aware hyperparameter optimization (HPO) in NQS calculations. Its workflow monitors energy trajectories, detects destructive optimization events, stops unstable calculations, modifies the learning-rate schedule, resumes optimization from safe checkpoints, and ranks candidates with an anomaly-aware score. We demonstrate the approach on a residual convolutional NQS for the square-lattice Heisenberg $J_1$-$J_2$ model, using architectures with parameter counts comparable to aCNN, a convolutional NQS architecture used here as a reference. The results show that NQS-Agent improves over the reported human-tuned aCNN baseline for the aCNN reference architecture and identifies a structurally distinct wide-and-shallow competitive candidate within the parameter-count-matched residual-CNN search space. These results show that the stability and recovery history of an optimization trajectory should be considered when assessing an NQS result. Health-aware HPO therefore provides a reproducible tuning protocol that goes beyond selecting a single lowest-energy calculation.

• The paper introduces NQS-Agent, an open-source framework that monitors NQS optimization trajectories, detects destructive events, resumes from safe checkpoints, and ranks candidates using energy and training health.
• The method improves on the human-tuned aCNN baseline at both tested couplings, reaching -0.495959(5) versus -0.495627(6) at J₂/J₁ = 0.50 and identifying competitive alternatives such as the wide-and-shallow A9 architecture.
• The results show that parameter count and FLOPs do not predict runtime or optimization quality, supporting reproducible HPO that reports architecture-dependent cost, recovery history, stability, and final energy together.

Motivation and problem statement

Neural-network quantum states (NQS) combined with variational Monte Carlo (VMC) have become a standard variational tool for strongly correlated lattice models, but their accuracy depends not only on the ansatz class but also on a large set of architecture- and optimizer-level hyperparameters: channel number, kernel size, residual-block count, learning rate, learning-rate schedule, regularization, seed, and optimization length. The paper argues that in practice these choices are made through manual expert intervention—inspecting energy curves, distinguishing transient from destructive instabilities, lowering learning rates, and resuming from saved states—and that such interventions do not scale to tens or hundreds of queued candidates. The central methodological claim is that an NQS candidate should be judged by its full optimization trajectory, including stability, checkpoint availability, and recovery history, rather than by the lowest energy reached at a fixed evaluation point.

The work introduces NQS-Agent, an open-source framework for health-aware hyperparameter optimization (HPO) of NQS, and demonstrates it on residual convolutional NQS for the square-lattice Heisenberg $J_1$-$J_2$ model, benchmarked against the human-tuned aCNN results (2606.30464).

Framework architecture

NQS-Agent is implemented on LangGraph and consists of three layers: an agent coordination layer with complementary roles (task interpretation, simulation design, simulation runner, data analysis, report generation), a persistent shared research state that stores candidate settings, energy logs, checkpoints, recovery decisions, and rankings, and a deterministic executable toolchain for HPO planning, input generation, health monitoring, checkpoint-based recovery, anomaly-aware scoring, and top-$k$ selection. A CLI/API entry exposes the same deterministic tools without the agent layer for reproducibility checks and batch execution; both routes share the same state and rules.

A key conceptual distinction is between a candidate experiment (physical problem, architecture, optimizer setting, seed) and an execution segment (one uninterrupted calculation). Several segments may belong to one candidate, so a promising candidate survives a destructive event if it can be resumed from a safe checkpoint rather than discarded or restarted.

Health diagnosis, recovery, and ranking

The monitor evaluates rolling-window statistics, tail variance and slope, non-finite values, and upward deviations from a trusted reference energy, which itself is updated only from windows whose normalized standard deviation does not exceed 5%. A destructive event—an unrecovered upward jump or any NaN/Inf value—triggers a reproducible sequence: stop the segment, select the latest safe checkpoint, modify the learning-rate schedule (halving shortly before the event, or capping at the pre-event value during the initial warm-up), and resume within the same candidate experiment. Recovery attempts are bounded, with the allowance increasing with target optimization length.

Ranking uses an anomaly-aware score combining the trusted median energy with penalties for tail deterioration, tail variance, positive late-time slope, transient events ($5\times10^{-4}$ each), destructive events ($2\times10^{-3}$ each), recovery attempts ($10^{-3}$ each), and missing full-resume checkpoints ($5\times10^{-4}$). Candidates are excluded outright when logs are invalid, destructive events remain unrecovered, no complete checkpoint exists, or the energy is discontinuous after continuation. The staged milestone-based allocation of longer runs to promoted candidates is related to Hyperband, ASHA, and BOHB, but differs in that promotion continues from a checkpoint preserving network and Adam optimizer state, and in that selection uses a task-specific top-$k$ rule incorporating trajectory health rather than a scalar objective. The workflow can additionally test the geometric mean $\sqrt{\eta_1\eta_2}$ of two leading learning rates as a targeted refinement.

Benchmark results

The benchmark uses $10\times10$ lattices at $J_2$0 and $J_2$1, where the human-tuned aCNN reference energies are $J_2$2 and $J_2$3 per site. The search space consists of nine representative residual CNNs (A1–A9) drawn from roughly 30 architectures whose parameter counts lie within $J_2$4 of the 6538-parameter aCNN reference (A5). Screening on $J_2$5 lattices retains A5, A6, and A9 for the $J_2$6 calculations.

Configuration	$J_2$7	$J_2$8
A5, LR=0.005, seed 8	$J_2$9	Not promoted
A5, LR=0.0013, seed 8	Not promoted	$k$0
A5, LR=0.002, seed 9	$k$1	$k$2
A5, LR=0.001265, seed 8	—	$k$3
A9, LR=0.008, seed 9	$k$4	Not promoted
A9, LR=0.020, seed 10	Failed after promotion	$k$5
A9, LR=0.0126, seed 12	Not promoted	$k$6
aCNN human-tuned baseline	$k$7	$k$8

Two claims follow. First, for the reference architecture A5, automated tuning with recovery improves over the human-tuned baseline at both couplings—for example, $k$9 versus $5\times10^{-4}$0 at $5\times10^{-4}$1. Second, A9, a structurally distinct wide-and-shallow network ($5\times10^{-4}$2, $5\times10^{-4}$3, $5\times10^{-4}$4), reaches energies below the baseline at both couplings despite passing through destructive fluctuations that were recovered via checkpoints. The authors are careful to state that because the study was not designed as a controlled head-to-head comparison across all seeds and settings, A9 should be regarded as a promising alternative rather than a definitive replacement for A5. The result nonetheless contradicts the assumption that deeper residual CNNs are always the most reliable route to accuracy.

Representative trajectories illustrate the mechanism concretely: an A5 run at learning rate 0.0013 exhibits one non-destructive fluctuation near step 650 and two destructive events (steps 840–978 and 2657–2788), after which checkpoint-based continuation allows the same candidate to proceed to a competitive final energy.

Architecture-dependent cost and speed

Parameter-count matching does not equalize computational cost. Forward FLOPs vary only within roughly $5\times10^{-4}$5–$5\times10^{-4}$6 the A5 value across the nine architectures, but stored activations span approximately $5\times10^{-4}$7–$5\times10^{-4}$8, and measured per-step speeds on an RTX 4090 GPU diverge sharply: A9 runs about $5\times10^{-4}$9 faster than A5 despite nearly identical FLOPs and stored activations, while the deepest candidate A4 runs at about $2\times10^{-3}$0 A5's speed. Faster architectures such as A3 and A7 were not selected by the energy- and stability-based HPO procedure, showing that execution speed and variational performance are distinct criteria. These observations support reporting architecture-dependent cost alongside energy in HPO studies, though the authors caution that the timing ratios may depend on the software environment and aCNN implementation details.

Limitations and open questions

The paper states several limitations explicitly. Architectures are matched by parameter count, not total computational cost, and the measured speed ratios are environment-dependent. Recovery rules are deterministic and deliberately conservative; more flexible agent-driven decisions for ambiguous fluctuations near trajectory ends remain untested and would need explicit safety constraints. The demonstrated search space covers only residual CNNs around the aCNN design on the Heisenberg $2\times10^{-3}$1-$2\times10^{-3}$2 model, so the generality of health-aware HPO across other NQS ansatz classes (RNNs, Transformers, graph networks) and other physical problems remains an open question. Additionally, whether agent-layer flexibility can outperform the fixed deterministic recovery heuristics without sacrificing auditability is left unresolved.

Conclusion

NQS-Agent converts expert manual interventions during NQS optimization into recorded, reproducible automated operations: trajectory monitoring, destructive-event detection, checkpoint-based recovery with schedule modification, and anomaly-aware ranking. On parameter-count-matched residual CNNs for the frustrated Heisenberg model, the protocol improves over the reported human-tuned aCNN energies for the reference architecture and identifies a wide-and-shallow competitor within the same parameter budget. The broader contribution is a shift in evaluation practice—from selecting the lowest final energy toward assessing optimization health jointly with energy—which makes large-scale automated NQS studies auditable and reproducible.