PaperlayerView exact revision on arXiv

Source-linked AI summary

NucleoNet: A Nuclear-Aware Determinant Ansatz for Multi-Geometry Neural Variational Monte Carlo

Qiao Gao, Yang Huang, Enze Hou, Yuzhi Liu, Jingrun Chen, Han Wang

arXiv:2610.12985v1physics.chem-phcs.LG

TL;DR

Neural wavefunctions are often optimized for one nuclear geometry, while mapping a potential energy surface requires accurate descriptions across geometries. NucleoNet jointly evolves electronic and nuclear representations in a shared determinant-wavefunction model, and one model describes multiple geometries in the tested molecular systems. Sparse-geometry training also generalizes to unseen N2 and LiCN geometries without further optimization.

  • Problem

    Separately optimizing neural wavefunctions for each geometry is costly when an accurate potential energy surface is needed.

    Source: §I · p. 1

  • Method

    NucleoNet processes electrons and nuclei as a unified particle set, jointly evolving their representations through a shared backbone to construct a determinant wavefunction.

    Source: §I · p. 2, §IV · p. 9

  • Results

    One jointly optimized model describes the tested N2, LiCN, ethane, and formamide geometry paths, while sparse-geometry models generalize to unseen N2 and LiCN geometries without further optimization.

    Source: §I · p. 2, §IV · p. 9

  • Takeaways & Limitations

    For the tested systems, a single parameter set can describe multiple geometries of a molecular system, including energy variation between training geometries on the tested N2 and LiCN coordinates.

    Source: §IV · p. 9

    Sparse-geometry generalization is shown along the tested one-dimensional N2 and LiCN coordinates, not the multi-nuclear ethane and formamide paths.

    Source: §IV · p. 9

Abstract

from arXiv · show

Neural-network variational Monte Carlo provides an accurate route to solving the electronic Schrödinger equation, but many neural wavefunctions are optimized for a single nuclear geometry, requiring separate calculations to map out a potential energy surface. We introduce NucleoNet, a nuclear-aware determinant ansatz in which electrons and nuclei form a unified particle set propagated through a shared electron--nucleus backbone that jointly evolves one-particle and two-particle representations. Nuclei carry learned representations that evolve alongside the electronic representations and condition orbital construction, while antisymmetry is imposed only by electronic determinants. On equilibrium small-molecule benchmarks, NucleoNet achieves accuracy comparable to established neural-wavefunction baselines with substantially fewer parameters. For each molecular system, a single model optimized jointly across geometries describes the N$_2$ dissociation curve, Li displacement in LiCN, and the ethane and formamide minimum-energy paths for bond dissociation. Sparse-geometry training generalizes to unseen geometries for N$_2$ and LiCN without further optimization, and N$_2$ ablations show that nuclear streams substantially improve this interpolation. These results show that a compact wavefunction with jointly evolved nuclear representations can describe the geometries of a molecular system with a single set of parameters.

I. INTRODUCTION

Mapping a molecular potential energy surface with separately optimized neural wavefunctions is costly and can produce noisy cross-geometry profiles. NucleoNet instead jointly evolves electronic and nuclear representations through a shared backbone to construct a determinant wavefunction.

I. INTRODUCTION

  • Separate optimization at each geometry is costly and can yield noisy potential energy surfaces because stochastic errors are uncorrelated across geometries.Neural VMC requires repeated Monte Carlo sampling, local-energy evaluation, wavefunction-gradient estimation, and stochastic parameter updates.Source: §I · p. 1
  • Earlier geometry-sharing approaches condition electronic wavefunctions on nuclear configurations but do not update nuclear representations in response to evolving electronic representations.The geometry pathway may use shared parameters, generated wavefunction parameters, nuclear representations, or orbital components.Source: §I · pp. 1–2

II. NUCLEONET

  • NucleoNet treats electrons and nuclei as a unified particle set whose representations are bidirectionally coupled through shared group summaries at every layer.The shared backbone provides electron-aware nuclear representations without a separate nuclear encoder, limiting geometry-conditioning parameter overhead.Source: §I · p. 2, §II · pp. 3–4
  • Nuclear representations influence the wavefunction both through backbone mixing into electronic representations and by generating geometry-dependent orbital-envelope parameters.Final electronic representations produce the Jastrow factor and orbital matrices, while nuclear representations generate envelope parameters.Source: §II · p. 2, Figure 1 · p. 3
  • NucleoNet uses a Slater–Jastrow multideterminant wavefunction, with fermionic antisymmetry enforced only by the electronic determinants.The electronic representations construct the orbital matrices, and the nuclear representations condition their envelopes.Source: §II · p. 2, §I · p. 2

III. RESULTS

The results are organized to test single-geometry accuracy, joint optimization across geometries, and generalization to geometries omitted from training. Ablations separate the roles of envelope conditioning and propagated nuclear streams.

III. RESULTS

  • The experiments first test whether NucleoNet retains competitive variational accuracy with a compact parameterization on equilibrium small molecules.The benchmark set comprises eight standard small molecules.Source: §III · pp. 4–5
  • A single model is optimized jointly over the N2 dissociation curve, Li displacement in LiCN, and ethane and formamide minimum-energy paths.The paths test whether one parameter set describes substantially different nuclear configurations along potential-energy surfaces.Source: §III · p. 4
  • Sparse-configuration optimization is evaluated on unseen geometries for N2 and LiCN, followed by ablations of envelope conditioning and nuclear streams.The ablations assess which component underlies the single model's behavior across geometries.Source: §III · pp. 4–5

A. Optimization protocol

For multi-geometry optimization, one shared parameter vector is trained over the selected nuclear configurations using their equally weighted mean variational energy. The optimization uses SPRING, a stochastic-reconfiguration-based optimizer.

B. Optimization protocol

  • A single parameter vector is shared across all training nuclear configurations, so each geometry contributes to optimization of the same wavefunction parameters.Electronic configurations are sampled from the corresponding geometry-dependent wavefunction.Source: A. Optimization protocol · p. 5
  • The multi-geometry objective is the equally weighted mean variational energy across training configurations.Each nuclear configuration contributes to the shared optimization.Source: A. Optimization protocol · p. 5
  • The variational optimization uses SPRING, a stochastic-reconfiguration optimizer that reuses recent-minibatch information without explicitly constructing and inverting the full SR matrix.SPRING is closely related to natural-gradient optimization in VMC.Source: A. Optimization protocol · p. 5

B. Single-geometry benchmarks

On eight equilibrium small-molecule benchmarks, NucleoNet reaches variational energies comparable to established neural-wavefunction baselines while using substantially fewer trainable parameters.

B. Single-geometry benchmarks

  • The benchmark compares optimized variational energies, for which lower energies indicate a better variational approximation up to statistical uncertainty.The eight-molecule benchmark follows the set and geometries used in previous neural VMC studies.Source: B. Single-geometry benchmarks · p. 5
  • 0.11–0.15 million trainable parameters give NucleoNet variational energies comparable to FermiNet, Psiformer, and LapNet across eight equilibrium small-molecule systems.Those baselines use several hundred thousand to multiple million parameters; Figure 2 plots energy error against parameter count.Source: Figure 2 · p. 5, B. Single-geometry benchmarks · p. 5

C. Nitrogen dissociation curve

Across N2, LiCN, ethane, and formamide, jointly optimized NucleoNet models describe energy profiles over multiple geometries. The reported results range from the lowest compared N2 curve MAE to close agreement with single-geometry or reference energies on the other paths.

C. Nitrogen dissociation curve

  • NucleoNet records a 1.07 mEh MAE on N2, the smallest among compared methods.The error profile is smooth, with largest deviations in the intermediate stretching region, and declines toward dissociation; one model is optimized over 32 uniformly spaced bond lengths from 1.6 to 6.4 a0.Source: C. Nitrogen dissociation curve · p. 6, Figure 3 · p. 6

D. Li displacement in LiCN

  • Across five LiCN geometries, joint optimization stays near the single-geometry baseline, with an MSD of +0.04 mEh.C and N remain fixed as Li moves along the molecular axis; the joint model lies below CCSD(T)/CBS and Orbformer at every geometry, so independently optimized NucleoNet energies provide the baseline.Source: D. Li displacement in LiCN · p. 6, Figure 4 · p. 6

E. Ethane and formamide dissociation pathways

  • Along the ethane and formamide cleavage paths, NucleoNet’s mean-centered profiles follow the references and its energies are lower than Orbformer’s at every image.Ethane stays within approximately 0.6 kcal/mol of the reference at every image (MSD −0.12 kcal/mol versus Orbformer’s 0.92); formamide has an MSD of 2.04 kcal/mol versus Orbformer’s 4.29.Source: E. Ethane and formamide dissociation… · pp. 6–7, Figure 5 · p. 7

F. Generalization from sparse geometry sets

Sparse-geometry NucleoNet models closely match full-geometry optimization on the tested N2 and LiCN coordinates. Ablations show that propagated nuclear streams particularly improve multi-geometry N2 performance.

F. Generalization from sparse geometry sets

G. Ablation studies

  • Fixing the envelope increases errors by 0.01 mEh for NH3, 0.11 mEh for single-geometry N2, and 0.21/0.23 mEh for N2 multi-geometry training/evaluation MAEs.Envelope conditioning provides a small, consistent improvement in these tested benchmarks.Source: G. Ablation studies · p. 8, Table I · p. 9
  • At fixed envelope, removing nuclear streams raises N2 training and evaluation MAEs by 0.71 and 38.63 mEh, with evaluation errors concentrated at geometries excluded from optimization.The ablation compares NucleoNet-FE-EB with NucleoNet-FE.Source: G. Ablation studies · p. 8, Figure 7 · p. 8, Table I · p. 9

IV. DISCUSSION

  • NucleoNet evolves each nucleus alongside electron streams through shared group summaries, while imposing antisymmetry only through electronic determinants.Source: §IV · p. 9
  • One parameter set describes each path for N2, LiCN, ethane, and formamide, spanning single-coordinate variation and multi-nuclear reaction paths.Source: §IV · p. 9

I. NETWORK IMPLEMENTATION DETAILS

NucleoNet builds particle and pair representations from type and geometry descriptors, then updates electron and nuclear streams through a shared backbone. This lets each species’ representations influence the other at every layer.

A. Multilayer Perceptron

  • Each nonlinear block applies a learned affine transformation and SiLU activation with residual mixing.The MLP is composed of multiple such blocks.Source: A. Multilayer Perceptron · p. 11

B. Particle and Pair Descriptors

  • Particle-type descriptors distinguish spin-up electrons, spin-down electrons, and nuclei using type, charge, and spin features.Pair descriptors combine the two particles’ type, charge, and spin features before the geometric features are added.Source: B. Particle and Pair Descriptors · p. 11

C. Nuclear-Aware Shared Backbone

  • When one- and two-particle representation dimensions differ, the one-particle representation is projected to the two-particle dimension before pair-gated messages are constructed.The source particle’s current one-particle representation and the pair representation both affect each ordered-pair message.Source: C. Nuclear-Aware Shared Backbone · pp. 11–12
  • NucleoNet updates electron and nuclear representations through shared group summaries, allowing each species’ stream to enter the other’s updates.At each layer, directed pair-gated messages depend on pair representations and the current representation of the source particle.Source: C. Nuclear-Aware Shared Backbone · p. 12

D. Determinant Construction

Final electronic representations construct orbital matrices, while final nuclear representations determine electron–nucleus envelope parameters. Determinants impose antisymmetry on electrons; the VMC optimization uses SPRING.

D. Determinant Construction

  • Orbital and envelope matrices are multiplied elementwise to form each spin-resolved block, which is stacked across electrons into a Slater matrix for each determinant.The nuclear representations determine envelope parameters, while electronic representations generate the orbital matrices.Source: D. Determinant Construction · p. 12
  • Antisymmetry arises because exchanging two same-spin electrons exchanges Slater-matrix rows and changes the determinant’s sign; nuclei are not included in that antisymmetrization.The rows are indexed only by electrons, and nuclear representations are not rows or columns of the Slater matrix.Source: D. Determinant Construction · p. 12

A. Optimizer and VMC Objective

  • All reported calculations use SPRING, which updates an approximate stochastic-reconfiguration direction using information from previous minibatches.The update is computed in sample space without constructing the full parameter-space SR matrix.Source: A. Optimizer and VMC Objective · p. 13
  • At each iteration, Monte Carlo samples define a subsampled stochastic-reconfiguration problem using centered log-wavefunction derivatives and local-energy fluctuations.The supplied passages identify these quantities as inputs to the subsampled problem but do not provide the full objective expression.Source: A. Optimizer and VMC Objective · p. 13

B. Optimizer and Sampling Hyperparameters

The supplement records optimizer, pretraining, architecture, and sampling choices, alongside the energy reporting conventions and benchmark data scope. Settings differ between equilibrium small molecules and multi-geometry experiments.

B. Optimizer and Sampling Hyperparameters

  • The reported calculations use SPRING and common optimizer hyperparameters unless otherwise stated, with the initial learning rate adjusted across benchmark settings.The calculations are reported as performed on a single NVIDIA A100 GPU.Source: B. Optimizer and Sampling… · p. 13, Table S1 · p. 13

C. Hartree–Fock Pretraining

  • Hartree–Fock pretraining precedes equilibrium small-molecule runs and the N2 and LiCN multi-geometry experiments, but is omitted for ethane and formamide.For ethane and formamide, pretraining produced slightly higher final VMC energies than direct VMC optimization.Source: C. Hartree–Fock Pretraining · p. 14

D. Architecture Hyperparameters

IV. MULTI-GEOMETRY BENCHMARK DETAILS

  • Multi-geometry optimization shares one parameter set across selected geometries and minimizes their average VMC energy.The supplement reports signed energy differences against reference calculations for these benchmarks, except N2, which uses mean absolute error against its empirical or fitted potential curve.Source: §IV · pp. 14–15

III. SMALL-MOLECULE BENCHMARK DATA

  • Small-molecule benchmark tables report total energies and parameter counts for eight systems, with FermiNet and Psiformer values from [3] and LapNet values from [4].The molecular geometries follow the FermiNet benchmark protocol; the table caption says bold values mark each system’s lowest reported energy.Source: §III · p. 14, Table S3 · p. 15

IV. MULTI-GEOMETRY BENCHMARK DETAILS

  • The multi-geometry supplement reports effective batch sizes, with nchain × Ngeom denoting chains per geometry times the number of geometries in the optimization.For ethane and formamide, the reference calculations are from Ref.; negative signed differences indicate lower optimized variational energies than those references.Source: §IV · p. 15, Table S4 · p. 15

A. N2 Dissociation Details

The supplementary experiments specify the geometry grids and reference calculations for joint and sparse-geometry tests. Sparse models are evaluated at geometries excluded from optimization without further wavefunction updates.

A. N2 Dissociation Details

  • A single NucleoNet model is jointly optimized and evaluated on 32 N2 bond lengths uniformly spaced from 1.6 to 6.4 a0.The reference is the MLR4(6, 8) potential of Ref., with its dissociation limit anchored to twice the atomic nitrogen energy per Ref..Source: A. N2 Dissociation Details · p. 15, Table S5 · p. 16

B. LiCN Benchmark Details

  • The LiCN benchmark varies Li along the molecular axis across five collinear configurations while C and N remain fixed.Li–C distances span approximately 3.6 to 4.7 a0; the study compares independently optimized geometries with one model jointly optimized across all five.Source: B. LiCN Benchmark Details · p. 16, Table S6 · p. 16

C. Ethane and Formamide Dissociation Details

  • Ethane and formamide use 20-image minimum-energy paths for homolytic C–H and N–H cleavage, respectively, with remaining nuclei allowed to relax.One NucleoNet model is jointly optimized across each path; reference energies are independently optimized single-geometry Psiformer VMC calculations from Ref..Source: C. Ethane and Formamide Dissociation… · pp. 16–17, Table S7 · p. 17

V. GENERALIZATION FROM SPARSE GEOMETRY SETS

A. N2 Generalization

  • For N2, a model trained on six bond lengths from 1.6 to 6.0 a0 is evaluated without parameter updates on a 32-point grid from 1.6 to 6.4 a0.The full-grid energies are reported in Table S10.Source: A. N2 Generalization · p. 17, Table S10 · p. 18

B. LiCN Generalization

  • For LiCN, a model trained only on the two endpoint configurations is evaluated on the three intermediate configurations without further wavefunction optimization.The resulting five-geometry profile is compared with a model jointly optimized across all five configurations.Source: B. LiCN Generalization · pp. 17–18

VI. ABLATION DETAILS

The ablations compare controlled changes to NucleoNet while keeping training, sampling, and optimization settings fixed. They isolate the roles of nuclear-conditioned envelope parameters and propagated nuclear representations.

VI. ABLATION DETAILS

  • The full architecture combines a shared electron–nucleus backbone with an envelope conditioned on final nuclear representations; the variants remove envelope conditioning, then restrict the backbone to electron streams.Other training, sampling, and optimization settings remain unchanged.Source: §VI · p. 18

A. Relative Metrics

  • Relative energy compares each ablation with full NucleoNet, while N2 sparse-geometry MAE changes are reported separately for training and evaluation geometries.Positive relative energy or ∆MAE indicates a higher energy or larger MAE, respectively, than the full model.Source: A. Relative Metrics · p. 19

B. Fixed-envelope ablation: NucleoNet-FE

  • NucleoNet-FE replaces nuclear-conditioned envelope parameters with geometry-independent trainable parameters while retaining nuclear-centered spatial dependence and the shared backbone.This comparison isolates conditioning the envelope parameters on learned nuclear representations.Source: B. Fixed-envelope ablation: NucleoNet-FE · p. 19

C. Electron-only-backbone ablation: NucleoNet-FE-EB

  • NucleoNet-FE-EB removes nuclear-representation propagation but retains static nuclear geometry through pooled initialization features and a nuclear-centered fixed envelope.It tests an electron-only backbone against NucleoNet-FE while retaining the same fixed envelope.Source: C. Electron-only-backbone ablation:… · pp. 19–20

D. Raw Ablation Data

The ablation analysis reports absolute variational energies in Hartree for the full model and its two variants across single-geometry and N2 sparse-geometry settings.

D. Raw Ablation Data

  • Tables S12–S14 report absolute variational energies in Hartree for NucleoNet, NucleoNet-FE, and NucleoNet-FE-EB, with uncertainty in parentheses.They cover NH3 and N2 single-geometry benchmarks, plus N2 sparse-geometry training and evaluation sets.Source: D. Raw Ablation Data · p. 20
  • Table S12 covers the NH3 and N2 single-geometry ablation benchmarks.Energies are reported with statistical uncertainty in the last reported digit.Source: Table S12 · p. 21
  • Table S13 gives N2 energies at the six geometries used to optimize sparse-geometry ablation models.Table S14 contains the corresponding N2 sparse-geometry evaluation-set results.Source: D. Raw Ablation Data · p. 20, Table S13 · p. 21
Loading 2610.12985v1…