Source-linked AI summary
Better, Faster Fermionic Neural Networks
James S. Spencer, David Pfau, Aleksandar Botev, W. M. C. Foulkes
TL;DR
FermiNet offers accurate neural-network wavefunctions but faces computational and scaling limitations on larger systems. This paper enlarges and simplifies the architecture, implements it in JAX, and evaluates it on larger atoms and challenging molecular reactions. The resulting methods reach chemical accuracy through argon, reduce large-system training cost by an order of magnitude, and obtain near-state-of-the-art reaction results.
Problem
FermiNet optimization requires substantial computational resources, and its accuracy diminishes slightly as system size grows.
Source: §1 · p. 2Method
The authors increase network size and combine JAX implementation with simplifications to improve FermiNet accuracy and training efficiency.
Source: §1 · p. 2, §3 · pp. 2–3Results
Takeaways & Limitations
The improvements extend accurate FermiNet calculations to larger atoms and make challenging molecular transitions computationally accessible.
Source: §3 · p. 4, §4 · p. 5FermiNet remains significantly slower to train than PauliNet, and the exact reason for this speed discrepancy remains unresolved.
Source: §4 · p. 5
Abstract
from arXiv · showhide
The Fermionic Neural Network (FermiNet) is a recently-developed neural network architecture that can be used as a wavefunction Ansatz for many-electron systems, and has already demonstrated high accuracy on small systems. Here we present several improvements to the FermiNet that allow us to set new records for speed and accuracy on challenging systems. We find that increasing the size of the network is sufficient to reach chemical accuracy on atoms as large as argon. Through a combination of implementing FermiNet in JAX and simplifying several parts of the network, we are able to reduce the number of GPU hours needed to train the FermiNet on large systems by an order of magnitude. This enables us to run the FermiNet on the challenging transition of bicyclobutane to butadiene and compare against the PauliNet on the automerization of cyclobutadiene, and we achieve results near the state of the art for both.
1 Introduction
The paper motivates accurate wavefunction solutions for many-electron systems and introduces variational quantum Monte Carlo with neural-network Ansätze, focusing on FermiNet’s scalability limitations.
- Accurate Schrödinger-equation solutions are central to quantum chemistry and condensed matter physics.Source: §1 · p. 1
- The time-independent equation seeks a wavefunction eigenfunction of the system Hamiltonian, with energy E as its eigenvalue.Source: §1 · p. 1
- Under the Born-Oppenheimer approximation, the lowest-energy solution is approximated by minimizing the Rayleigh quotient over parameterized unnormalized wavefunctions.Source: §1 · p. 1
- Monte Carlo evaluation of this variational energy using samples from pθ(x) ∝ ψθ^2(x) is known as variational quantum Monte Carlo (VMC).Source: §1 · p. 1
- Deep neural networks provide flexible wavefunction Ansätze, while FermiNet achieved higher absolute accuracy on many small atomic and molecular systems than other neural-network Ansätze.Source: §1 · pp. 1–2
- The original FermiNet’s computational demands and declining accuracy with system size motivate improvements in scaling, accuracy, and optimization.Source: §1 · p. 2
2 Improving FermiNet Accuracy
The authors improve FermiNet accuracy on second-row atoms by enlarging the network and extending training, reaching chemical accuracy for phosphorus through argon except sulphur.
- Increasing both determinant count and one-electron-stream width reached chemical accuracy for phosphorus through argon except sulphur.Chemical accuracy is defined as 1 kcal/mol to the exact energy.Source: §2 · p. 2
- Training used 20 MCMC steps between weight updates and 300,000–400,000 updates to reduce noise, improve equilibration, and ensure convergence.Source: §2 · p. 2
3 Scaling and Simplifying FermiNets
The authors simplify and reimplement FermiNet to reduce computational cost while preserving or improving accuracy, enabling challenging molecular benchmarks.
3 Scaling and Simplifying FermiNets
- JAX raises large-system GPU utilization from ~60% to ~90% and reduces bicyclobutane’s required GPUs from 16 to 4, yielding 6x greater efficiency.The reduction applies to a batch size of 4096.Source: §3 · p. 3
- 40% lower computational overhead leaves bicyclobutane performance unaffected when full covariance is replaced by a single isotropic parameter.Source: §3 · p. 3
- 10x faster training and 8 mEh lower overall energy result from combining JAX implementation with network simplifications.This result uses the original network width and determinant count, without the larger settings explored earlier.Source: §3 · pp. 3–4
Bicyclobutane to 1,3-butadiene transition
- FermiNet matches DMC energy differences for all bicyclobutane-to-butadiene states within chemical accuracy without system-specific tuning.It also closely matches CC(t;3), except for slightly lower energy on the disrotatory pathway; Table 3 reports benchmark energies relative to bicyclobutane with zero-point corrections.Source: Table 3 · p. 3, §3 · p. 4
Cyclobutadiene automerization and comparison to the PauliNet
- FermiNet reaches ~70 mEh lower absolute energy than PauliNet on cyclobutadiene, while their relative barriers are comparable: 10.3±0.1 versus 9.9±0.6 kcal/mol.Both relative barriers lie at the high end of experimentally observed values, and the FermiNet agrees well with the highest MR-CC result.Source: §3 · p. 4, Figure 2 · p. 4
4 Conclusions
The paper combines engineering and architectural simplification to accelerate FermiNet training while increasing accuracy through larger networks, extending accurate applications to difficult atomic and molecular systems.
4 Conclusions
- FermiNet reaches chemical accuracy across the second-row atoms and produces chemically accurate relative energies for the bicyclobutane-to-butadiene transition.This includes the disrotatory transition state, where CCSD(T) fails dramatically.Source: §4 · p. 5
- FermiNet achieves better absolute energies than PauliNet on cyclobutadiene but comparable relative energies.Source: §4 · p. 5
- FermiNet remains significantly slower to train than PauliNet, with the discrepancy potentially reflecting learned versus built-in cusp conditions or an accuracy–speed trade-off.The exact reason for the training-speed difference remains unresolved.Source: §4 · p. 5
Broader Impact
The work could lead to adoption of new computational techniques by chemistry and materials science communities, but its present scope remains limited to already well-understood systems.
- The methods could be adopted by chemistry and materials science communities.Source: Broader Impact · p. 5
- The present work is too small for cutting-edge chemistry and has so far been used only on already well-understood systems.Source: Broader Impact · p. 5