Case Studies

Learning Quantum Mechanical Wave Functions with Baskerville

What is this research about?

Professor Matthew Foulkes is a theoretical materials physicist interested in all types of everyday matter — from molecules to metals, insulators, semiconductors, magnets, and superconductors. His research ranges from the practical — studying the formation and properties of protective oxide scales on high-performance alloys — to the theoretical, such as superconductivity in idealised model systems. The common thread is that most projects begin from the quantum mechanical many-particle Schrödinger equation, which serves as a grand unified theory for materials physics.

Gaining the ability to solve the Schrödinger equation accurately and efficiently would be transformative, allowing materials scientists, physicists, chemists, and biochemists to participate in the simulation-based revolution that has already overtaken much of engineering. This would greatly reduce the cost of discovering and developing new materials, chemicals, catalysts, and medicines. In practice, the many-particle Schrödinger equation is extremely difficult to solve for any but the smallest systems, and researchers almost always rely on approximate methods such as density functional theory.

Four years ago, working with a small team from Google DeepMind, Professor Foulkes introduced FermiNet — a new neural-network-based approach to solving the many-particle Schrödinger equation. Similar approaches had been tried for model systems, but this was the first successful application to real molecules and solids. The FermiNet neural network takes any set of particle positions as input and returns the corresponding value of the many-particle wavefunction. The network weights are optimised using the variational principle, which states that the best approximation to the ground-state wavefunction can be found by minimising the energy expectation value. No external data are required.

The computational cost is high, but FermiNet and its successor the Psiformer are much more accurate than density functional theory, scale better with system size than comparably accurate quantum chemistry approaches, and run naturally on massively parallel GPU-based supercomputers such as Baskerville.

Recent work, much of which used Baskerville, has shown that neural wave functions can describe strongly correlated systems beyond the reach of DFT, and that they can “discover” quantum phase transitions — such as Wigner crystallisation and the onset of superconductivity in the electron gas — without human guidance. They appear to be particularly effective at describing unusual systems where conventional methods struggle. Recent work on the binding of positrons to molecules, for example, produced some of the most accurate computations of binding energies and annihilation rates available.

Why was Baskerville chosen?

The codes, developed in close collaboration with researchers at Google DeepMind, were designed from the start to run on GPU supercomputers, leveraging their vast low-level parallelism via the JAX library, the XLA accelerated linear algebra compiler, and NVIDIA’s CUDA programming model. Large runs require 40 GB A100 GPUs or better.

GPU resources available to university-based researchers in the UK are limited, and were even more limited when the group applied for their first Baskerville grant through EPSRC’s Access to High Performance Computing mechanism. Baskerville was one of only two A100 machines available at the time. Although there was limited choice, Baskerville proved to be powerful, stable, and well run. The support received was always helpful and efficient, and the availability of large-memory 80 GB A100 GPUs was a significant benefit. Much of the group’s research over the past few years would have been impossible without Baskerville.

Selected publications

  • D. Pfau, J.S. Spencer, A.G.D.G. Matthews, and W.M.C. Foulkes, Phys. Rev. Research 2, 033429 (2020).
  • I. von Glehn, J.S. Spencer, and D. Pfau, arXiv:2211.13672 (2022).
  • G. Cassella, H. Sutterud, S. Azadi, N.D. Drummond, D. Pfau, J.S. Spencer, and W.M.C. Foulkes, Phys. Rev. Lett. 130, 036401 (2023). (Editors’ Suggestion)
  • W.T. Lou, H. Sutterud, G. Cassella, W.M.C. Foulkes, J. Knolle, D. Pfau, and J.S. Spencer, Phys. Rev. X 14, 021030 (2024).
  • G. Cassella, W.M.C. Foulkes, D. Pfau, and J.S. Spencer, Nat. Commun. 15, 5214 (2024).