🔬 Research: FermiNet: Fermionic Neural Network for quantum mechanics - paper/code

[8bitmp3]

“If you think you understand quantum mechanics, you don’t understand quantum mechanics.” — Richard Feynman

Just wanted to highlight another paper with JAX code called Ab initio solution of the many-electron Schrödinger equation with deep neural networks by @dpfau, @jsspencer, @alexggmatthews (DeepMind), and Matthew Foulkes.

Additional JAX-based libraries include Chex (utils) and Optax (for processing and optimizing grads).

📄 Paper: https://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.2.033429
💻 JAX code: https://github.com/deepmind/ferminet/tree/jax
🔗 DeepMind's blog post: FermiNet: Quantum Physics and Chemistry from First Principles

TL;DR (blog post)

In an article recently published in Physical Review Research, we show how deep learning can help solve the fundamental equations of quantum mechanics for real-world systems. Not only is this an important fundamental scientific question, but it also could lead to practical uses in the future, allowing researchers to prototype new materials and chemical syntheses in silico before trying to make them in the lab.
...
We’ve developed a new neural network architecture, the Fermionic Neural Network or FermiNet, which is well-suited to modeling the quantum state of large collections of electrons, the fundamental building blocks of chemical bonds. The FermiNet was the first demonstration of deep learning for computing the energy of atoms and molecules from first principles that was accurate enough to be useful, and it remains the most accurate neural network method to date.

Deep neural networks can often be far more efficient at representing complex functions than linear combinations of basis functions. In the FermiNet, this is achieved by making each function going into the determinant a function of all electrons (1). This goes far beyond methods that just use one- and two-electron functions. The FermiNet has a separate stream of information for each electron. Without any interaction between these streams, the network would be no more expressive than a conventional Slater determinant. To go beyond this, we average together information from across all streams at each layer of the network, and pass this information to each stream at the next layer. That way, these streams have the right symmetry properties to create an antisymmetric function. This is similar to how graph neural networks aggregate information at each layer. Unlike the Slater determinants, FermiNets are universal function approximators, at least in the limit where the neural network layers become wide enough. That means that, if we can train these networks correctly, they should be able to fit the nearly-exact solution to the Schrödinger equation.

The paper also says the authors "greatly improve the accuracy of the Slater-Jastrow-backflow variational quantum Monte Carlo (VMC) method by using a neural network we dub the Fermionic Neural Network, or FermiNet, as a more flexible Ansatz".