Accelerating Chemical Potential Calculations with Minimal Normalizing Flows

Chemical potentials are among the most important properties that can be obtained from a molecular simulation since they define many technologically relevant collective properties such as solubilities and activity coefficients. The chemical potential of a species in solution is typically obtained by computing the free energy change of adding that species into a bulk system, a calculation typically very expensive for systems of technological interest such as electrolytes, due to the lack of phase space overlap between "not-inserted″ and "inserted″ states. Recently, normalizing flows have been introduced as a way to accelerate free energy computations by learning a bijective function to map the configuration space of one Boltzmann distribution onto another. Currently, these trainable mappings are constructed to be as expressive as possible, theoretically having the ability to create a perfect mapping between states. This expressivity makes them difficult to train, limiting their ability to be generated "on-the-fly″ for any new free energy estimation challenge, and in practice these mappings have shown only modest performance improvements for liquid systems. We address these issues by introducing the concept of a "minimal″ normalizing flow (MNF). This is a trainable bijective mapping that is intentionally limited in expressivity and instead applies low-dimensional, physically informed transformations. Useful MNFs can be trained in approximately 1 min of GPU time due to their simplicity and our introduction of a novel loss function based on the Bhattacharyya distance as an alternative to the Kullback-Leibler divergence. We show how calculations of chemical potentials of pure and binary Lennard-Jones particle systems can be accelerated by at least an order of magnitude with a simple radial mapping. For a more complex illustration of the approach, we apply a two-dimensional (radial and orientational) mapping to the solvation of sodium and chloride ions in water, showing that MNFs can increase the effective sample size by as much as 3 times for charging free energy calculations and as much as 8 times for calculating free energy changes due to force field perturbations. This provides the foundation for the development of physically informed normalizing flows that can accelerate complex free energy calculations while retaining low training costs.