Normalizing flows for atomic solids

A flow model is a probability distribution q defined as the pushforward of an analytically tractable base distribution b through a flexible diffeomorphism f, typically parameterized by neural networks [14, 15]. Independent samples from q can be generated in parallel, by first sampling z from b and taking $x = f(z)$. The probability density of a sample can be obtained using a change of variables,

$$\begin{equation} q(x) = b(z)\left|{\det J_f(z)}\right|^{-1}, \end{equation} \tag{ 2 }$$

where J_f is the Jacobian of f.

We can train q to approximate p by minimizing a loss function that quantifies the discrepancy between q and p. In this work, we use the following Kullback–Leibler divergence as the loss function

$$\begin{align} D\!\left({q}\|{p}\right) & = \left\langle{\ln{q(x)} - \ln{p(x)}}\right\rangle_{q} \nonumber\\ & = \left\langle{\ln{q(x)} + \beta U(x)}\right\rangle_{q} + \ln{Z} \nonumber \\ & = \left\langle{\ln{b(z)} - \ln{\left|{\det J_f(z)}\right|} + \beta U(f(z))}\right\rangle_{b} + \ln{Z}. \end{align} \tag{ 3 }$$

Since $\ln{Z}$ is a constant with respect to the parameters of q, it can be ignored during optimization. The expectation in the final expression can be estimated using samples from b, so $D\!\left({q}\|{p}\right)$ can be minimized using stochastic gradient-based methods. This loss function is appealing because it does not require samples from p or knowledge of Z, and can be optimized solely using evaluations of the energy U.

The systems considered in this work are crystalline solids consisting of N indistinguishable atoms at constant volume and temperature. The crystal is assumed to be contained in a 3-dimensional box with edge lengths $L_1, L_2, L_3$ and periodic boundary conditions. A configuration x is an N-tuple $(x_1, \ldots, x_N)$, where $x_n = (x_{n1}, x_{n2}, x_{n3})$ are the coordinates of the nth atom and $x_\mathrm{ni} \in [0, L_i/\sigma]$ with σ being a characteristic length scale of the system (here the particle diameter). By expressing x in reduced units, Z and all probability densities become dimensionless.

The potentials used in this work are invariant to global translations (with respect to periodic boundary conditions) and arbitrary atom permutations. Both of these symmetries are incorporated into the model architecture; that is, we design the base distribution b and the diffeomorphism f such that the density function q is invariant to translations and atom permutations, as well as compatible with periodic boundary conditions. We explain how this is achieved in the following paragraphs.

Our base distribution b is constructed as a lattice with N sites that have been randomly perturbed and permuted as illustrated in figure 1. Starting from a lattice $z_o = (z_{o1}, \ldots, z_{oN})$, a configuration $z = (z_{1}, \ldots, z_{N})$ is generated by independently adding spherically-truncated Gaussian noise to each lattice site, followed by a random permutation of all atoms. The truncation is chosen such that no two neighbouring atoms can swap lattice sites, so that after the permutation, all atoms can be traced back to the site they originated from. This construction yields a base distribution that can be trivially sampled from, and has a permutation-invariant probability density function that can be evaluated exactly.

Zoom In Zoom Out Reset image size

Our diffeomorphism f is implemented as a sequence of K invertible functions composed such that $f = f_K \circ \cdots \circ f_1$. Each function f_k is parameterized by a separate neural network, whose parameters are optimized by the loss in equation (3). We implement the functions f_k using an improved version of the model proposed in [21]. In this model, each f_k transforms element-wise either one or two coordinates of all atoms as a function of all remaining coordinates. The transformation of each coordinate is implemented using circular rational-quadratic splines [27], which ensures that the transformation is nonlinear, invertible and obeys periodic boundary conditions. The spline parameters are computed as a function of the remaining coordinates using multiple layers of self-attention, a neural-network module commonly used in language modelling [28]. Wirnsberger et al [21] showed that a diffeomorphism f parameterized this way is equivariant to atom permutations, which means that permuting the input of f has the same effect as permuting the output of f. As also shown by [25, 29], the combination of a permutation-invariant base distribution with a permutation-equivariant diffeomorphism yields a permutation-invariant distribution q, as desired. More implementation details are provided in the supplementary material (available online at stacks.iop.org/MLST/3/025009/mmedia).

Finally, we incorporate the translational symmetry into the above architecture as follows. We fix the coordinates of an arbitrary reference atom (say x₁), and use the flow model to generate the remaining N − 1 atom coordinates as described above. Then, we globally translate the atoms uniformly at random (under periodic boundary conditions), so that the reference atom can end up anywhere in the box with equal probability. Since the index of the reference atom and its original position are known and fixed, we can reverse this operation in order to obtain the probability density of an arbitrary configuration. This procedure yields a translation invariant probability density function that can be calculated exactly.

A key feature of our model architecture is that we can target specific crystal structures by encoding them into the base distribution. For example, if we are interested in modelling the hexagonal phase of a crystal, we can choose the lattice of the base distribution to be hexagonal. Empirically, we find that, after training, the flow model becomes a sampler for the (metastable) crystal state that we encode in the base distribution, and does not sample configurations from other states. Thus, by choosing the base lattice accordingly, we can guide the model towards the state of interest, without changing the energy function or using ground-truth samples for guidance.

To assess the quality of our trained models, we compare against MD. We ran NVT simulations of the target systems using the simulation package LAMMPS [35]. Figure 2 shows the energy histogram and radial distribution function (RDF) of the Lennard-Jones FCC crystal and of cubic ice, as computed by samples from the base distribution, the trained model, and MD. The RDF g(r) is the ratio of the average number density at a distance r of an arbitrary reference atom and the average number density in an ideal gas at the same overall density [2]. By construction, the base distribution captures the locations of the peaks in the RDF correctly, but its energy histogram is far off compared to the MD result. The energy histograms and the RDFs computed from model samples, however, are nearly indistinguishable from MD in both systems. Importantly, no unbiasing or re-weighing was necessary to obtain this quality of agreement. The results for hexagonal ice show a similar level of agreement (see supplementary material). These results demonstrate that the mapping f successfully transforms the base distribution into an accurate sampler.

Zoom In Zoom Out Reset image size

To demonstrate that the trained model becomes a sampler of the (metastable) crystal structure that we have encoded into the base distribution, we computed histograms of the averaged local bond-order parameter q₆ [36] that was designed to discriminate between different phases (see figure 3). We can see that the two histograms are well separated and agree with MD, showing that the model does indeed become an accurate sampler of only the crystal state encoded in the base distribution. If that were not the case, we could have enforced it by adding a biasing potential to the energy, as commonly done in nucleation studies that employ umbrella sampling [37].

Zoom In Zoom Out Reset image size

Next, we use the trained flow models to estimate the Helmholtz free energy F for various system sizes, which is given by [2]

$$\begin{equation} F = -\beta^{-1}\left(\ln{Z}-\ln{N!}\right). \end{equation} \tag{ 4 }$$

We note that the thermal de Broglie wavelength does not appear in the above expression as we set it to σ, following [38], and absorb it into Z by expressing x in reduced units. We then estimate $\ln{Z}$ by first defining a generalized work function [19]

$$\begin{equation} \beta \Phi(x) = \beta U(x) + \ln{q(x)}. \end{equation} \tag{ 5 }$$

The average work value $\left\langle{\beta \Phi(x)}\right\rangle_{q}$ is also our training objective in equation (3). We then harness the trained flow to draw a set of samples $\{x^{(m)}\}_{m = 1}^M$ from q and estimate $\ln{Z}$ via the targeted free energy perturbation estimator [19]

$$\begin{equation} \ln{Z} = \ln{\left\langle{\exp(-\beta \Phi(x))}\right\rangle_{q}} \approx \ln{\frac{1}{M}\sum_{m = 1}^M\exp\!\left(-\beta\Phi(x^{(m)})\right)}. \end{equation} \tag{ 6 }$$

The above estimation method, referred to as learned free energy perturbation (LFEP) in combination with a learned model [21], is appealing, because it does not require samples from p for either training the model or for evaluating the estimator.

Although the approximation in equation (6) becomes exact in the limit of an infinite sample size, to obtain accurate results for a finite sample set we require sufficient agreement between the proposal and the target distributions [19, 21]. Since $\beta \Phi$ quantifies the pointwise difference between $\ln{q}$ and $\ln{p}$ (up to an additive constant), the distribution of generalized work values is a good metric for assessing the quality of the flow for free energy estimation.

Figure 4 compares distributions of work values computed for the base distribution and for the fully trained model. From the non-negativity of the Kullback–Leibler divergence in equation (3) it follows that $\left\langle{\beta \Phi(x)}\right\rangle_{q} \geqslant -\ln{Z}$, with equality if and only if q and p are equal. Therefore, the gap between the average work value and $-\ln{Z}$ (which we aim to estimate) quantifies how accurately the model q approximates the Boltzmann distribution p; for a perfect model, we would expect to see a delta distribution located at $-\ln{Z}$ [19]. The work values obtained with the trained model are indeed sharply peaked near our MBAR estimate of $-\widehat{\ln{Z}}$. On a qualitative level, this shows a clear benefit of LFEP [21] over the original FEP estimator corresponding to f being the identity map, which failed to converge on this problem.

Zoom In Zoom Out Reset image size

Using the same trained model as in LFEP, we can also employ a learned version of the bidirectional BAR estimator (LBAR) [21, 39]. LBAR uses samples from both the base distribution b and from the target p and is known to be the minimum variance estimator for any asymptotically unbiased method [40]. The downside compared to LFEP is that an additional MD simulation needs to be performed to obtain samples from p for computing the estimator (but not for training the model).

To verify the correctness of the flow-based free energy estimates, we require an accurate baseline method. The Einstein crystal (and molecule) method [38, 41] and lattice-switch Monte Carlo (LSMC) [42] are common choices for computing solid free energies with different trade-offs. While the former is conceptually simple, we found it challenging to optimize it to high precision, which is consistent with previously reported results [30, 38]. The latter is known to be accurate but provides only free energy differences between two compatible lattices rather than absolute free energies. We therefore tested MBAR as an alternative estimator on this problem and found it yields sufficiently accurate absolute free energy estimates to serve as a reference.

A quantitative comparison of free energy estimates for both systems and different system sizes is shown in table 1 (see also supplementary material). For LFEP we use S samples from the base distribution b; for LBAR we use S samples from b, plus another S samples from the target p obtained by MD; for MBAR we use a sufficiently large number of intermediate states between b and p to obtain good accuracy, which we sample using MD (see caption of table 1 for exact numbers). Overall, we find excellent agreement of the learned estimators with MBAR for both LJ and ice. The LBAR estimates exhibit lower statistical uncertainties than LFEP across the board, with error bars on the order of $10^{-5} k_\mathrm{B} T$ per particle. We find it remarkable, however, that LFEP can yield comparable accuracy in most cases without access to MD samples for training or estimation, and without the need for defining intermediate states. Finally, we compute the Helmholtz free energy difference between cubic and hexagonal ice for 216 particles by subtracting the two LFEP estimates in table 1. This yields a value of $12.4(2)~\mathrm{J\,mol^{-1}}$ which is in good agreement with the reported Gibbs free energy difference of $11.2(2)~\mathrm{J\,mol^{-1}}$ obtained with LSMC simulations at atmospheric pressure [31].