FermiNet: learning fermionic wavefunctions with variational Monte Carlo

A neural network that represents a wavefunction rather than a label

FermiNet is a neural network for learning highly accurate ground state wavefunctions of atoms and molecules using a variational Monte Carlo approach.

That description hides an unusual thing about the project: the output is not a prediction in the usual sense. The network is trained to output a wavefunction, and that wavefunction is then sampled with Monte Carlo walkers to produce energies and other observables. The network parameters are variational parameters, and minimising the energy is the training objective.

The repository contains the implementation of the algorithm and experiments first described in Ab-Initio Solution of the Many-Electron Schroedinger Equation with Deep Neural Networks, by David Pfau, James S. Spencer, Alex G de G Matthews and W.M.C. Foulkes, Phys. Rev. Research 2, 033429 (2020), along with subsequent research and developments.

So the audience is computational chemistry and quantum machine learning researchers. If you want to compute the energy of a molecule for a paper, this is the wrong tool and the package that replaced Hartree-Fock and DFT is the right one. If you want to study the method itself, to extend it to new systems, or to reproduce the numbers in the paper and the follow-ups, the code is the artefact.

The project is Apache-2.0 licensed, and the package metadata names DeepMind as the author with a version of 0.2.

The JAX rewrite moved the original TensorFlow code to a branch

The README opens with a warning that is worth reading before anything else: this is a research-level release of a JAX implementation and is under active development, and the original TensorFlow implementation can be found in the tf branch.

That single sentence tells you most of what you need to know about the state of the codebase. The current code is a port, not the original. It was rewritten on JAX, which is the right choice for a project that needs custom autodiff through a Monte Carlo sampling loop and automatic multi-device parallelism, but a rewrite also means the two implementations are not interchangeable.

If you have existing TensorFlow FermiNet code, or if a paper you are reproducing used the TensorFlow version, the tf branch is where the matching implementation lives. If you are starting fresh, the main branch is the one to use.

Two other signals point the same way. There are no GitHub releases, so there is no stable version tag to pin, and the last push to the repository was on 2026-10-15. The repository root also carries a CONTRIBUTING.md, a pylintrc, a pytype.cfg and a .flake8 file, which is the tooling profile of a project where code quality is enforced by static analysis rather than by a published release contract.

Installing into a virtualenv, with a jaxlib version you must match

Installation is a standard editable install, and the README recommends doing it inside a virtual environment.

virtualenv ~/venv/ferminet
source ~/venv/ferminet/bin/activate
pip install -e .

The setup script declares the dependencies, and the list tells you what this project really is. Alongside the expected jax and jaxlib, there is kfac_jax from DeepMind, folx from Microsoft, optax for optimisation, ml-collections for configuration, chex for test assertions, pyscf for chemistry, pyblock, absl-py, h5py, numpy, scipy and pandas.

Two of those deserve comment. kfac_jax is the Kronecker factored optimiser that makes training a wavefunction network practical, and pyscf means the code can take a molecule defined in a well known quantum chemistry package rather than only hand-built coordinates.

GPU support is handled separately, and the README is explicit about the constraint: the jaxlib version must correspond to the existing CUDA installation you wish to use.

pip install --upgrade jax jaxlib==0.1.57+cuda110 -f
https://storage.googleapis.com/jax-releases/jax_releases.html

Note that the command pins a jaxlib build for CUDA 11.0, which is an older generation. If your machine has a newer CUDA, that exact pin will not be the one you want, and the JAX documentation is linked for the current matrix.

Tests are run with pytest, and the testing extra installs what they need.

pip install -e '.[testing]'
python -m pytest

Configuration is a ConfigDict, and the command line flags are the real interface

Configuration uses the ConfigDict from ml_collections, and the detail that matters is that it is a Python object you edit, not a text file of key value pairs.

The packaged example scripts live under ferminet/configs/, and the README warns that most of them are there for testing, so a production calculation may need additional settings. The reference for every available setting is ferminet/base_config.py.

The command line is the primary interface, because ConfigDict supports flag overrides directly.

ferminet --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100

There is an equivalent module invocation for anyone who does not want the console entry point.

python3 ferminet/main.py --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100

One inconsistency is worth flagging before you copy these commands. The text after the examples says the run trains on the Li atom using a batch size of 1024 MCMC configurations and 100 iterations of pretraining, while the flag in the command sets the batch size to 256. The command is the authority, but if your results differ from a number quoted in a paper, check which batch size was actually used before suspecting the code.

The idea of a walker is worth holding on to. The batch size is the number of Monte Carlo configurations sampled in parallel, so it is not a conventional mini-batch, and changing it changes the variance of the energy estimate rather than just the memory footprint.

New systems are configured in Python, not declared in a schema

Other systems are set up by creating a new config file or by writing a custom training script. The README shows an H2 molecule defined in a config file with a get_config function.

from ferminet import base_config
from ferminet.utils import system

# Settings in a config files are loaded by executing the the get_config
# function.
def get_config():
  # Get default options.
  cfg = base_config.default()
  # Set up molecule
  cfg.system.electrons = (1,1)
  cfg.system.molecule = [system.Atom('H', (0, 0, -1)), system.Atom('H', (0, 0, 1))]

  # Set training hyperparameters
  cfg.batch_size = 256
  cfg.pretrain.iterations = 100

  return cfg

The file is then run like any other config.

ferminet --config /path/to/h2_config.py

The important detail is the electron count. cfg.system.electrons is a pair of alpha and beta electrons, and the second number matters because spin is part of the state the network represents.

If you prefer to define the molecule with PySCF, the code accepts a PySCF Molecule directly.

from pyscf import gto
mol = gto.Mole()
mol.build(
    atom = 'H  0 0 1; H 0 0 -1',
    basis = 'sto-3g', unit='bohr')

That molecule is then assigned to cfg.system.pyscf_mol and passed to the training entry point, which is the path to take when you want the geometry and basis from a package that already knows how to handle them.

Writing a custom script instead of a config file is equally supported, and it gives you control over logging, since the absl logging handler can be pointed at standard output for progress output during a script run.

Inference is the training command with the optimizer switched off

After training, you re-run the same command with the optimiser disabled, using the flag --config.optim.optimizer 'none'.

The reasoning is sound and worth understanding. Training produces a set of parameters, and one run of sampling at those parameters gives a noisy estimate of the energy. To get a low variance estimate you need to accumulate over many time steps with the parameters frozen, which is what inference does.

There is one operational requirement: either the value of cfg.log.save_path must be the same, or the value of cfg.log.restore_path must be set to the value of cfg.log.save_path from the original training run. In other words, inference has to find the checkpoint that training wrote, and the README does not leave you to guess which of the two fields to touch.

The second use of inference is collecting observables that were not tracked during training. Spin magnitude, dipole moments and density matrices can be turned on at inference time by adding --config.observables.s2, --config.observables.dipole and --config.observables.density to the command line, if they are not already set to true in the config file.

That separation is genuinely useful in practice. A training run is expensive, and being able to compute additional physical properties afterwards from the same parameters, without retraining, changes what you can do with a single run.

The one hard scaling requirement sits alongside this: to train on larger atoms and molecules with large batch sizes, multi-GPU parallelisation is essential, supported through JAX pmap, with multiple GPUs automatically detected and used if available.

Excited states come in two flavours, and the paper configs are included

Excited state properties can be calculated with either the Natural Excited States for VMC algorithm, or an ensemble penalty method.

Enabling the calculation of k states is a single setting, cfg.system.states=k. NES-VMC is the default, and switching to the ensemble penalty method means adding cfg.optim.objective='vmc_overlap' to the config.

The two approaches differ in how much you have to configure them. NES-VMC has no parameters to set. The ensemble penalty method has a free choice of weights on the energies and the overlap penalty, set in cfg.optim.overlap, and if the energy weights are not set for the states in the config they are automatically set to 1/k for state k.

The README states a preference, noting that NES-VMC is generally more accurate than the ensemble penalty method but that both are included for completeness. That is a claim from the authors about their own two implementations, and it is the kind of statement you should treat as a hypothesis to check on your own systems rather than a settled result.

Practical help comes with the two methods. Config files for all experiments from the paper that introduced NES-VMC are in the folder configs/excited, and the README notes that all experiments can be tested on small systems. Those configs are the fastest route to a working excited state calculation, since they carry the settings that were actually used for the published results.

Editorial conclusion

FermiNet suits a computational chemistry group that wants to reproduce the DeepMind results, extend them to new systems, or work on excited states, since it ships the algorithm, example configs for many experiments and both the NES-VMC and ensemble penalty routes to excited states. It does not suit a chemist who wants a drop-in DFT or wavefunction package for routine calculations, and it does not suit a machine without a GPU for anything beyond toy systems, because the README is explicit that multi-GPU parallelisation is essential for larger atoms and molecules with large batch sizes. Before starting, read the research release warning at the top of the README, confirm your jaxlib version matches your installed CUDA version, and check that the package version you are tracking is the JAX one on main rather than the older TensorFlow implementation that lives on the tf branch.

Frequently asked questions

Official sources

1. google-deepmind/ferminet on GitHub
2. Issues
3. License: Apache-2.0
4. README