Optimizing Free Energy Estimation with Machine Learning

Peter Wirnsberger
Senior Research Scientist
DeepMind

Computer-based free energy estimation has been an active field of research for decades, with many successful applications in physics, materials science and biology. Free energy perturbation (FEP) [1] is a bedrock technique for estimation of free energy differences. A requirement for fast and reliable convergence of the estimator, however, is that the respective distributions share a large overlap in configuration space. An elegant strategy to address this problem is to increase overlap using configuration space mappings. This approach, known as Targeted Free Energy Perturbation (TFEP) [2], has the appealing property that it can achieve immediate convergence when combined with a perfect mapping. Defining a good high-dimensional mapping, however, is a challenging task. 

In this talk, I will present our recent work [3] in which we turn TFEP into a machine learning problem. In our approach, we represent the mapping by a deep neural network whose parameters are optimized so as to maximize overlap. We test our targeted estimators on a prototypical solvation system for which we generate training data using molecular dynamics simulations. Our neural network architecture respects two important symmetries of the system, namely periodic boundary conditions and permutational symmetry of identical particles. We observe that our technique leads to an improved accuracy of the free energy estimates compared to baselines, without requiring any additional data for training the network.

References
[1] R. W. Zwanzig, J. Chem. Phys. 22, 1420 (1954).
[2] C. Jarzynski, Phys. Rev. E 65, 046122 (2002).
[3] This research by P. Wirnsberger, A.J. Ballard, G. Papamakarios, S. Abercrombie, S. Racanière, A. Pritzel, D. Jimenez Rezende, and C. Blundell has been first published in J. Chem. Phys. 153 , 144112 (2020), with the permission of AIP Publishing. PW and AJB contributed equally to this work.

Speaker Bio
Peter is a Senior Research Scientist at DeepMind. His research interests include developing new machine learning algorithms as well as applying them to problems that lie at the intersection of machine learning and physics. Peter obtained his PhD in Theoretical Chemistry from the University of Cambridge, where he worked with Daan Frenkel on quantifying polarization effects in molecular systems far from thermal equilibrium.