Estimating Gibbs free energies via isobaric-isothermal flows

We consider a system of N atoms in D dimensions, confined in a fixed simulation box (typically orthorhombic). Let $x_n = (x_{n,1}, \ldots, x_{n,D})$ be the position coordinates of the nth particle, and $x = (x_1, \ldots, x_N)$ be the configuration of the system. We assume that the atoms are subject to a potential energy U(x) and the system is in contact with a heat bath at fixed temperature T with which it can exchange energy. A standard result in statistical mechanics is that the equilibrium distribution of such a system (known as the canonical or NVT ensemble) is given by the Boltzmann distribution:

$$\begin{equation} p(x) = \frac{1}{Z_V}\exp[-\beta U(x)], \end{equation} \tag{ 1 }$$

where $Z_V = \int \mathrm{d}x \exp[-\beta U(x)]$ is the canonical partition function and $\beta = 1/k_\mathrm{B} T$ is the inverse temperature ($k_\mathrm{B}$ is Boltzmann's constant). The canonical partition function is directly related to the Helmholtz free energy F via

$$\begin{equation} F = -\beta^{-1}\ln{Z_V}. \end{equation} \tag{ 2 }$$

In practice, one is interested in estimating (i) expectations under the Boltzmann distribution ${\mathbb{E}}_{p(x)} \left [ {A(x)} \right]$ of physical observables A, and (ii) the partition function Z_V and hence the free energy F. These problems are computationally challenging due to the high dimensionality of x and the fact that p(x) can only be evaluated point-wise and up to a constant. In what follows, we focus on a method that takes advantage of recent advances in ML and deep generative modeling, specifically normalizing flows.

Normalizing flows are deep generative models for parameterizing flexible probability distributions. A normalizing flow parameterizes a complex distribution $q_\theta(x)$ (with parameters θ) as the pushforward of a (typically fixed) simple base distribution q(u) via an invertible and differentiable function f_θ (typically parameterized with deep neural networks). In other words, a sample x from $q_\theta(x)$ is drawn by first sampling u from q(u) and transforming it with f_θ :

$$\begin{equation} x \sim q_\theta(x) \quad\Leftrightarrow\quad x = f_\theta(u), u\sim q(u). \end{equation} \tag{ 3 }$$

Thanks to the invertibility and differentiability of f_θ , the probability density of x can be computed exactly with a change of variables:

$$\begin{equation} q_\theta(x) = q(u)\left|{\det \frac{\partial f_\theta}{\partial u}}\right|^{-1}, \end{equation} \tag{ 4 }$$

where $\frac{\partial f_\theta}{\partial u}$ is the Jacobian of f_θ . There is a wide variety of implementations of invertible and differentiable functions f_θ using deep neural networks that are flexible enough to parameterize complex distributions $q_\theta(x)$ but whose Jacobian determinant $\det \frac{\partial f_\theta}{\partial u}$ remains efficient to compute, enabling exact density evaluation; for a thorough review, see [4, 15].

Given a target distribution that can be evaluated point-wise and up to a constant, such as the Boltzmann distribution in equation (1), the objective is to train a normalizing-flow model to approximate it. Let the intractable target distribution take the general form $p(x) = \frac{1}{Z}\exp[g(x)]$ and let $q_\theta(x)$ be a flow model as defined in equation (3). We train the flow to approximate the target by minimizing a suitably chosen divergence between p(x) and $q_\theta(x)$ with respect to θ. A common choice is the reverse Kullback–Leibler divergence:

$$\begin{align} D\!\left({q_\theta(x)}\,\|\,{p(x)}\right) & = {\mathbb{E}}_{q_\theta(x)} \left [ {\ln{q_\theta(x)} - \ln{p(x)}} \right]\kern65.5pt \end{align} \tag{ 5 }$$

$$\begin{align} & = {\mathbb{E}}_{q_\theta(x)} \left [ {\ln{q_\theta(x)} - g(x)} \right] + \ln{Z} \end{align} \tag{ 6 }$$

$$\begin{align} & \quad\quad= {\mathbb{E}}_{q(u)} \left [ {\ln{q(u)} - \ln{\left|{\det \frac{\partial f_\theta}{\partial u}}\right|} - g(f_\theta(u))} \right] + \ln{Z}, \end{align} \tag{ 7 }$$

where in the last step we reparameterized the expectation with respect to the base distribution q(u). Using the final expression, the gradient of the KL divergence with respect to θ is

$$\begin{equation} \nabla_\theta D\!\left({q_\theta(x)}\,\|\,{p(x)}\right) = {\mathbb{E}}_{q(u)} \left [ {- \nabla_\theta\ln{\left|{\det \frac{\partial f_\theta}{\partial u}}\right|} - \nabla_\theta g(f_\theta(u))} \right]. \end{equation} \tag{ 8 }$$

Using samples from q(u), one can readily obtain an unbiased estimate of the above gradient, which can be used to minimize the KL divergence with any stochastic gradient optimization method.

After we train a flow model $q_\theta(x)$ to approximate the target distribution p(x), samples from $q_\theta(x)$ can be used in a Monte-Carlo fashion to estimate expectations ${\mathbb{E}}_{p(x)} \left [ {A(x)} \right]$. Since $q_\theta(x)$ does not exactly equal p(x), direct Monte-Carlo estimates with samples from $q_\theta(x)$ will be biased; to avoid that, the probability density of samples can be used to remove the bias. Assuming $q_\theta(x) \gt 0$ whenever $p(x) \gt 0$, the simplest way to remove bias is with importance weighting: using i.i.d. samples $(x^{(1)}, \ldots, x^{(S)})$ from $q_\theta(x)$, we can form the asymptotically unbiased estimate

$$\begin{equation} {\mathbb{E}}_{p(x)} \left [ {A(x)} \right] \approx \frac{\sum_{s = 1}^S w(x^{(s)})A(x^{(s)})}{\sum_{s = 1}^S w(x^{(s)})} \quad\text{where}\quad w(x) : = \frac{\exp[g(x)]}{q_\theta(x)}. \end{equation} \tag{ 9 }$$

Similarly, the normalizing constant $Z = \int \mathrm{d}x \exp[g(x)]$ can be unbiasedly estimated as

$$\begin{equation} Z \approx \frac{1}{S}\sum_{s = 1}^S w(x^{(s)}). \end{equation} \tag{ 10 }$$

A disadvantage of importance-weighting is that it results in an estimator whose statistical performance is lower than that of a standard Monte-Carlo average of S independent samples. To quantify this reduction in statistical performance, the effective sample size (ESS) is often used, estimated by [16]

$$\begin{equation} {\mathrm{ESS} \approx \frac{\left[\sum_{s = 1}^S w(x^{(s)})\right]^2}{\sum_{s = 1}^S \left[w(x^{(s)})\right]^2}.} \end{equation} \tag{ 11 }$$

The ESS is a number between 1 and S (often reported as a percentage of S), and roughly measures how many independent samples would result in a Monte-Carlo estimator of the same statistical performance.

In statistical mechanics, the estimation method in equation (10) is widely known as free energy perturbation (FEP) [17]; when combined with a configuration space map it is referred to as targeted FEP [14] and, in combination with a learned map or a flow model it is referred to as learned FEP (LFEP) [6, 7].

In general, the simpler the estimation method, the closer the flow model must approximate the target for the statistical efficiency of the estimator to be within practical constraints. More sophisticated, but also computationally more expensive, multi-step estimation methods are available and can be utilized as learned estimators [6, 18–21] if higher accuracy is desirable.

In the isobaric-isothermal ensemble, the external pressure P is held fixed while the shape of the simulation box, and hence its volume V, fluctuates. Therefore, the shape parameters of the simulation box should be treated as additional degrees of freedom that co-vary with the configuration x. In what follows, we will assume that the shape of the simulation box is controlled by a set of parameters h and we will formulate a joint probability density $p(x, h)$.

A standard result in statistical mechanics is that the equilibrium distribution of the (x, V) system is given by replacing the potential energy U with the enthalpy U + PV in the Boltzmann distribution in equation (1), which gives:

$$\begin{equation} p(x, V) = \frac{1}{Z_P} \exp\Big[-\beta(U(x, V) + P V)\Big], \end{equation} \tag{ 12 }$$

where $Z_P = \int \mathrm{d}x \,\mathrm{d}V \exp[-\beta (U(x, V) + PV)] = \int \mathrm{d}V \exp[-\beta PV] Z_V$ is the partition function of the isobaric-isothermal ensemble. We note that the potential energy generally depends on the box shape, hence we include an explicit dependence on V. Analogously to the NVT case, the partition function Z_P is related to the Gibbs free energy G via

$$\begin{equation} G = -\beta^{-1}\ln{\,Z_P}. \end{equation} \tag{ 13 }$$

The above distribution is directly applicable to simulation boxes whose only shape parameter is V, for example a simulation box with isotropic scaling. However, it is not directly applicable to simulation boxes with more than one shape parameter without further assumptions of how the additional shape parameters are distributed.

In this work, we are interested in the case of a triclinic simulation box in D dimensions. The full flexibility to adjust the box shape is, for example, important when studying solids, to ensure that they can relax to their equilibrium shape [12, 13]. To describe the triclinic box shape, we let h be a D × D matrix with positive determinant that transforms the hypercube $[0, 1]^D$ into the simulation box. For D = 3, such a simulation box has nine shape parameters controlling the following nine degrees of freedom: scaling each of the three side lengths, varying each of the three internal angles, and three rigid rotations (see figure 1). Assuming that all simulation boxes with the same volume V are a priori equally likely, Martyna et al [22] derive the following distribution for the (x, h) system (equation 2.15 in their paper):

$$\begin{equation} p(x, h) = \frac{1}{Z_P} \exp\Big[-\beta (U(x, h) + PV) + (1-D)\ln{V}\Big], \end{equation} \tag{ 14 }$$

where $V = \det{h}$ is the volume of the simulation box.

Zoom In Zoom Out Reset image size

Since the systems we consider are globally rotationally symmetric, it is desirable to eliminate the $D(D-1) / 2$ rigid rotations from the shape parameters h. This can be done by restricting h to be upper-triangular ($h_{ij} = 0$ for i > j) and requiring its diagonal elements h_ii to be positive. In this case, Martyna et al [22] derive the following distribution (equation 2.35 in their paper):

$$\begin{equation} p(x, h) = \frac{1}{Z_P} \exp\Bigg[-\beta (U(x, h) + P V) + (1-D)\ln{V} + \sum_{i = 1}^D(i-1)\ln{h_{ii}}\Bigg], \end{equation} \tag{ 15 }$$

where $V = \prod_{i = 1}^D h_{ii}$ is the volume of the simulation box.

A practical issue with the above distribution is that the range of atom positions $x = (x_1, \ldots, x_N)$ is not fixed, but varies with the shape parameters h. To avoid this complication, we define a fixed reference box with parameters h₀ (where we choose h₀ to be a diagonal D × D matrix with positive diagonal elements) and work with reference coordinates $s = (s_1, \ldots, s_N)$ that are defined with respect to the reference box. Concretely, we make the following variable changes:

• (i)  
  $h = M h_0$, where M is a D × D upper-triangular matrix with positive diagonal elements, and
• (ii)  
  $x_n = M s_n$ for $n = 1, \ldots, N$.

In addition, we log-transform the diagonal elements of M, so that we work with unconstrained shape parameters and that we ensure the diagonal elements M_ii are always positive. That is, we define a D × D upper-triangular matrix m with unconstrained upper-triangular elements, and make the variable change:

• (iii)  
  $M_{ii} = \exp(m_{ii})$ for $i = 1, \ldots, D$, and
• (iv)  
  $M_{ij} = m_{ij}$ for i ≠ j.

The above transformations change the variables of interest from (x, h) to (s, m). As shown in the supplementary material, this change of variables leads to the following probability density for (s, m), which is the target distribution we use in our experiments:

$$\begin{align} p(s, m) = \frac{1}{Z_P} \exp\Bigg[&-\beta (U(x, h) + P V) + (N + 2 - D)\ln{V} \nonumber \\ &- (N+1-D) \ln{V_0} + \sum_{i = 1}^D(i-1)\ln{h_{ii}}\Bigg], \end{align} \tag{ 16 }$$

where $V = \prod_{i = 1}^D \exp(m_{ii})h_{0,i}$ is the volume of the simulation box and $V_0 = \prod_{i = 1}^D h_{0,i}$ is the volume of the reference box. The term $(N+1-D) \ln{V_0}$ is a constant with respect to (s, m), so it can be absorbed into the normalizing constant.

In this section, we describe the architecture of a model $q_\theta(s, m)$ that is suitable for approximating the target in equation (16). Our primary design choice is to decompose the model as:

$$\begin{equation} q_\theta(s, m) = q_{{\theta_1}}(m) q_{{\theta_2}}(s \,|\, m), \end{equation} \tag{ 17 }$$

where $q_{{\theta_1}}(m)$ is a flow model over the shape parameters, and $q_{{\theta_2}}(s \,|\, m)$ is a shape-conditional flow model over the configuration. The trainable parameters of the full model are $\theta = (\theta_1, \theta_2)$, where θ₁ are the trainable parameters of $q_{\theta_1}(m)$ and θ₂ are the trainable parameters of $q_{\theta_2}(s \,|\, m)$. In the following paragraphs, we describe each model in detail. The overall model architecture is illustrated in figure 2.

Zoom In Zoom Out Reset image size

3.2.1. Flow over shape parameters

3.2.2. Shape-conditional flow over configuration

We first justify our use of a Gaussian-affine flow by verifying that for our systems of interest the true distribution of shape parameters is statistically indistinguishable from a six-dimensional Gaussian. To this end, we compare histograms of the six shape parameters sampled from the flow model with their counterparts sampled during MD. The results for the 512-particle cubic system are shown in figure 3 and demonstrate excellent agreement of model and MD. Results for hexagonal ice show the same quality of agreement (see supplementary material) and results for all smaller systems are omitted for the sake of brevity, as training with fewer particles is in general much simpler due to the reduced number of degrees of freedom.

Zoom In Zoom Out Reset image size

To further support our assumption that a Gaussian with full covariance matrix is sufficient to model the shape distribution accurately for ice, we estimated the maximum mean discrepancy (MMD) [28] between the MD distribution of shape parameters and a full-covariance Gaussian fit. We observed that MMD estimates are close to zero (zero MMD implies equality of distributions) and that they are distributed similarly to reference MMD estimates between two identical six-dimensional Gaussians (see figure 4).

Zoom In Zoom Out Reset image size

We next compare the distributions of energies and instantaneous pressure values for both cubic and hexagonal 512-particle systems in figure 5 and observe again good overall agreement between model and MD for both quantities. Table 1 shows the average energy and pressure for various system sizes (64, 216 and 512 particles) for both cubic and hexagonal ice, as estimated (a) directly by model samples (biased), (b) by re-weighted model samples (unbiased), and (c) by MD. We also show the ESS of importance sampling in each case as a percentage of the actual sample size. The results suggest that the bias due to the model is generally small, and that both observables are accurately estimated. We can further see that importance sampling is able to remove small biases, although at the expense of increasing the error bars, especially for the largest system size (512 particles). Finally, we see that the ESS is reasonably high for 64 and 216 particles, but becomes small for 512 particles, leading to the pronounced increase in estimation uncertainty mentioned previously.

Zoom In Zoom Out Reset image size

Table 1. Biased (BE) and unbiased importance-sampling (IS) estimates of mean energy and pressure for system sizes of N = 64, 216 and 512 particles, along with estimates of the ESS expressed as a percentage of the actual sample size. MD estimates were computed with 100 seeds and 10 000 samples each. Model estimates comprise 16 seeds for each phase for 64 particles, 16 and 14 seeds for cubic and hexagonal, respectively, for 216 particles, and 11 seeds for 512 particles, with 512 000 samples per seed in all cases.

	Cubic		Hexagonal
	Energy (kcal mol−1)	Pressure (atm)	Energy (kcal mol−1)	Pressure (atm)
N = 64	ESS = $53.7\%$		ESS = $60.3\%$
BE	$-750.61 \pm 0.02\hphantom{0}$	$0.88 \pm 1.10$	$-750.50 \pm 0.02\hphantom{0}$	$0.28 \pm 1.98$
IS	$-750.456 \pm 0.005$	$1.15 \pm 0.92$	$-750.351 \pm 0.005$	$0.29 \pm 1.98$
MD	$-750.451 \pm 0.008$	$0.85 \pm 0.63$	$-750.352 \pm 0.009$	$1.30 \pm 0.73$
N = 216	ESS = $6.8\%$		ESS = $13.6\%$
BE	$-2535.08 \pm 0.05$	$1.58 \pm 1.26$	$-2534.69 \pm 0.04$	$1.03 \pm 1.46$
IS	$-2534.50 \pm 0.11$	$2.43 \pm 2.67$	$-2534.20 \pm 0.01$	$-0.01 \pm 0.80$
MD	$-2534.54 \pm 0.02$	$0.93 \pm 0.39$	$-2534.22 \pm 0.02$	$0.80 \pm 0.42$
N = 512	ESS = $0.17\%$		ESS = $0.15\%$
BE	$-6009.56 \pm 0.17\hphantom{0}$	$1.57 \pm 0.91$	$-6009.07 \pm 0.11\hphantom{0}$	$0.93 \pm 1.14$
IS	$-6008.84 \pm 0.48\hphantom{0}$	$-1.63 \pm 16.1$	$-6007.6 \pm 1.1\hphantom{00}$	$-23.6 \pm 38.8$
MD	$-6008.776 \pm 0.027$	$0.98 \pm 0.28$	$-6008.011 \pm 0.024$	$0.93 \pm 0.24$

To further assess the quality of the trained model, we computed the radial distribution function and compare the results for both phases to MD in figure 6. We find the NPT-flow results to be almost indistinguishable from MD under this metric, suggesting that the model captures pairwise, structural properties correctly.

Zoom In Zoom Out Reset image size

We finally compare Gibbs free energy estimates obtained with our trained models with MD estimates for all three system sizes. For the MD baseline estimates, we first borrow the Helmholtz free energy estimates published in [11] and promote these to Gibbs free energy estimates following the method proposed by Cheng and Ceriotti [29]. The Gibbs free energy is related to the Helmholtz free energy by

$$\begin{equation} G(P, T) = F(h_\text{ref}, T) + P \det(h_\text{ref}) + k_\text{B} T \ln p (h_\text{ref} \,|\, P, T), \end{equation} \tag{ 18 }$$

where h_ref defines the fixed simulation cell for which we know the Helmholtz free energy estimate, and $\ln p (h_\text{ref} \,|\, P, T)$ denotes the log-density of the reference box at fixed pressure and temperature. To obtain an estimate of this quantity, we first fit a six-dimensional Gaussian to the shape data sampled in an NPT-simulation, and then evaluate the log-density of the fit on the reference box analytically.

Table 2 contains the Gibbs free energy differences predicted by our models, using the LFEP estimator in equation (10), with the MD baseline estimates computed using equation (18). We find excellent agreement of both methods across all system sizes, verifying that our NPT-flow can resolve even the fine free energy difference between cubic and hexagonal ice accurately.