Source-linked AI summary
Targeted free energy estimation via learned mappings
Peter Wirnsberger, Andrew J. Ballard, George Papamakarios, Stuart Abercrombie, Sébastien Racanière, Alexander Pritzel, Danilo Jimenez Rezende, Charles Blundell
TL;DR
Free-energy perturbation estimates require sufficient overlap between distributions, and designing useful targeted mappings is difficult. This work learns mappings with neural networks that respect atomic symmetries, and reports variance reduction over the respective estimator baselines on a solvation system.
Problem
Free-energy perturbation convergence depends on sufficient overlap, while crafting a tractable targeted mapping for realistic systems is challenging.
Source: §I · p. 1, §VII · p. 7Method
The paper parameterizes targeted mappings with neural networks optimized to increase overlap and designs the architecture to respect permutation and periodic symmetries.
Source: §I · p. 2Results
Both LFEP and LBAR significantly reduce variance compared with their respective baselines in the tested solvation system.
Source: §VII · p. 7Takeaways & Limitations
The results highlight the potential of learned mappings for targeted free-energy estimation, with staging and mapping possibly combined in complex systems.
Source: §VII · p. 7The unidirectional LLFEP loss does not encourage mass coverage, and its KL minimum does not yield the lowest bias.
Source: Appendix D · p. 9
Abstract
from arXiv · showhide
Free energy perturbation (FEP) was proposed by Zwanzig more than six decades ago as a method to estimate free energy differences, and has since inspired a huge body of related methods that use it as an integral building block. Being an importance sampling based estimator, however, FEP suffers from a severe limitation: the requirement of sufficient overlap between distributions. One strategy to mitigate this problem, called Targeted Free Energy Perturbation, uses a high-dimensional mapping in configuration space to increase overlap of the underlying distributions. Despite its potential, this method has attracted only limited attention due to the formidable challenge of formulating a tractable mapping. Here, we cast Targeted FEP as a machine learning problem in which the mapping is parameterized as a neural network that is optimized so as to increase overlap. We develop a new model architecture that respects permutational and periodic symmetries often encountered in atomistic simulations and test our method on a fully-periodic solvation system. We demonstrate that our method leads to a substantial variance reduction in free energy estimates when compared against baselines, without requiring any additional data.
I. INTRODUCTION
Free energy estimation is broadly useful but finite-sample FEP convergence depends strongly on overlap between the sampled and target distributions. Targeted mappings can improve overlap, but designing them is difficult; this work learns them with neural networks and builds in relevant atomic symmetries.
I. INTRODUCTION
- Computer-based free-energy estimation is used across fields including drug discovery, condensed matter physics, materials science, and structural biology.The paper describes it as a valuable alternative where experimental methods and analytic theory cannot reach many systems.Source: §I · p. 1
- FEP convergence depends strongly on overlap because the dominant contributions come from samples typical under the target distribution, which become rarer as overlap decreases.FEP is an importance-sampling estimator, with samples from state A serving as the proposal distribution.Source: §I · p. 1
- Multistage approaches improve convergence by requiring sufficient overlap between neighboring states, but need samples from multiple states and a suitable order parameter.The number and discretization of intermediate stages are not clear a priori.Source: §I · p. 1
- Targeted FEP uses an invertible configuration-space mapping to transport samples toward a distribution with greater overlap, potentially improving convergence.The bidirectional extension applies the mapping and its inverse to samples from A and B.Source: §I · p. 1
- The challenge is constructing a tractable mapping without relying on physical intuition, motivating a learned mapping optimized to increase overlap.The neural-network mapping supports both unidirectional LFEP and bidirectional LBAR estimation.Source: §I · pp. 1–2
- The neural mapping is equivariant to permutations of identical particles and respects periodic boundary conditions by construction.These symmetries are common in atomic systems.Source: §I · p. 2
II. THEORETICAL BACKGROUND
Targeted estimators use invertible mappings to transform samples from thermodynamic states and estimate their free-energy difference. The bidirectional formulation applies BAR to forward and reverse mapped work, while finite-sample convergence depends on the mapping.
II. THEORETICAL BACKGROUND
- An invertible mapping transports configurations sampled from A to a mapped state A′; the reverse process maps samples from B using the inverse.The mapped distributions are related to their source distributions through the change-of-variable formula.Source: §II · p. 2
- Figure 1 contrasts multistage estimation, which sums differences across neighboring states, with targeted mapping, which seeks greater overlap between mapped and target distributions.The figure depicts regions of appreciable probability mass in configuration space.Source: Figure 1 · p. 2
- The generalized FEP identity recovers the free-energy difference from mapped work and reduces to ordinary FEP when the mapping is the identity.The identity holds for any invertible mapping.Source: §II · p. 2
- In the bidirectional setting, the fluctuation theorem relates forward and reverse processes, and BAR can estimate the free-energy difference from both.BAR is minimum-variance among asymptotically unbiased methods, so it is generally preferred when samples from both states are available.Source: §II · pp. 2–3
- The targeted estimators are consistent for any invertible mapping with infinitely many samples, but finite-sample convergence depends on the mapping choice.LBAR optimizes work values before combining them to predict the free-energy difference.Source: §II · p. 3
III. TRAINING OBJECTIVE
The training objective is motivated by overlap between mapped image distributions and their targets, since greater overlap improves estimator convergence. The paper uses KL-based objectives for unidirectional and bidirectional training, with the bidirectional loss encouraging mass coverage in both directions.
III. TRAINING OBJECTIVE
- When image and target distributions coincide, the work distributions collapse to the free-energy difference and estimator convergence is immediate.The same overlap argument applies to both generalized FEP and the bidirectional estimator.Source: §III · p. 3
- The training objective uses KL divergence between image and target distributions to judge mapping quality, with constants such as ΔF and β omitted during optimization.The KL cannot be accurately estimated unless ΔF is known, but ΔF and β are constants for optimization.Source: §III · p. 3
- Unidirectional training uses only the forward process, whereas bidirectional training uses both forward and reverse processes.The paper assigns a separate loss function to each regime.Source: §III · p. 3
- When samples from both states are available, the bidirectional loss is preferable because it encourages both mapped distributions to be mass-covering.Mass coverage is important for good performance of importance-sampling estimators.Source: §III · p. 3
IV. EXPERIMENTAL SETUP
The experiment estimates the free-energy change as a repulsive solute grows in a periodic bath of 125 solvent particles. The authors aim to learn a tractable mapping rather than hand-craft one, while accounting for the system’s symmetries.
IV. EXPERIMENTAL SETUP
- The task is to calculate the free-energy change as a soft repulsive solute grows from radius RA to RB in a bath of 125 identical solvent particles.The particles occupy a cubic periodic box.Source: §IV · p. 3, Figure 2 · p. 4
- Uniformly compressing the solvent space is not directly applicable to soft solutes because the transformed and target densities have different support, violating invertibility.The earlier mapping gave convergence gains for a hard solute.Source: §IV · pp. 3–4
- The proposed framework automates mapping design instead of relying on hand-crafted physical intuition, while requiring compatibility with the physical system’s symmetries.The authors identify symmetry compatibility as crucial for learning an effective mapping.Source: §IV · p. 4
A. Energy
The system uses periodic boundary conditions, and its total energy is decomposed into pairwise and particle-centered contributions. The interaction terms use radially symmetric potentials with a specified cutoff treatment.
A. Energy
- Periodic boundary conditions place the N-particle system on a 3N-dimensional torus, with each particle coordinate in [−L,L].The total energy is decomposed into pairwise contributions.Source: A. Energy · p. 4
- The displacement r_ij is computed using the shortest signed distance in each periodic dimension.The displacement is formed by subtracting an integer multiple of 2L from the coordinate difference.Source: A. Energy · p. 4
- The energy uses radially symmetric Lennard-Jones and Weeks–Chandler–Andersen potentials; Lennard-Jones interactions are truncated at L and shifted to zero there.Source: A. Energy · p. 4
B. Training data
Training and evaluation data come from Langevin molecular-dynamics simulations, with decorrelated frames split into independent datasets. The training-set size was selected to give the baseline BAR estimator a reasonably accurate free-energy difference.
B. Training data
- Langevin molecular dynamics generated 5×10^4 decorrelated frames per simulation, across 20 independent trajectories for each thermodynamic state.The trajectories began from randomly initialized configurations and momenta.Source: B. Training data · p. 4
- The authors constructed 10 datasets by concatenating and shuffling configurations, then splitting each into 9×10^4 training and 1×10^4 test samples.Source: B. Training data · p. 4
- The training-set size was chosen so the baseline BAR estimator would yield a reasonably accurate free-energy difference on the training data.Source: B. Training data · p. 4
C. Symmetries
The system has periodic-boundary and permutation symmetries, which the mapping is designed to respect by construction. The model handles the remaining cube symmetries through training-data augmentation.
C. Symmetries
- Because torus translations leave the energy invariant, the solute can be fixed at the box origin without changing the ratio ZB/ZA.The remaining discrete rotations and reflections form the 48-element octahedral group.Source: C. Symmetries · p. 4
- The mapping is constructed to respect periodic boundary conditions and permutation invariance of the identical solvent particles.The system’s periodic geometry is a torus.Source: C. Symmetries · p. 4
- The architecture does not guarantee octahedral symmetry, so training data are augmented with random elements of the 48-element group.The model is trained to learn these symmetries from the transformed data.Source: C. Symmetries · pp. 4–5
- Augmenting over all 125! solvent-particle permutations is infeasible in reasonable training time.The model instead respects permutation invariance by construction.Source: C. Symmetries · pp. 4–5
V. MODEL
The mapping is implemented as a composition of invertible neural-network layers, designed to permit efficient Jacobian calculations while respecting the system’s periodic and permutation symmetries.
V. MODEL
- The mapping M is a normalizing flow composed of K invertible layers, combining efficient Jacobian determinants with increased flexibility through composition.The first three design requirements are satisfied by normalizing flows; bijectivity is preserved when layers are composed.Source: §V · p. 5
- Each coupling layer leaves a subset of coordinates unchanged and transforms the rest element-wise using parameters computed from the invariant coordinates.Different coordinate partitions across layers let each coordinate be transformed as a function of the others.Source: §V · p. 5
- Circular splines implement the coordinate transforms so that the mapping respects periodic boundary conditions.The spline is required to map the periodic domain consistently, including identifying its boundary points.Source: §V · pp. 5–6
- The network uses periodic feature transformations for invariant coordinates, with sine and cosine features that preserve their information.These features make the parameters periodic functions of the invariant coordinates.Source: §V · p. 6
- A permutation-equivariant transformer produces the coupling-layer parameters, making the mapping’s Jacobian invariant to particle permutations.Equivariance ensures that permuting particles permutes parameter outputs in the same way.Source: §V · p. 6
VI. RESULTS
In the bidirectional solvation test, learned mapping increased work-distribution overlap and reduced LBAR’s estimate variance relative to BAR. The training and test losses also showed overfitting, motivating early stopping.
VI. RESULTS
- The test loss reaches a minimum before training loss stops decreasing, so the authors use the minimum-test-loss parameters for evaluation.The figure places the test-loss minimum at around 2.1×10^5 training steps; the authors report that its precise location does not appreciably affect bidirectional estimates.Source: §VI · p. 6, Figure 4 · p. 6
- The mapped forward and reverse work distributions shifted toward one another, increasing their overlap and reducing their variance.The crossing of the distributions corresponds to the free-energy difference; Fig. 5a compares mapped and unmapped work distributions.Source: §VI · p. 6, Figure 5 · p. 7
- LBAR’s full-batch standard deviation was approximately 19% of BAR’s, while its estimate variation overlapped the MBAR error estimates despite MBAR using 7.5 times more total samples.Training and evaluation used the same dataset, demonstrating data-efficient mapping learning in this experiment.Source: §VI · p. 7, Figure 5 · p. 7
VII. DISCUSSION
The discussion presents learned mappings as a way to make targeted free-energy estimation tractable and reports variance reduction on a periodic solvation prototype. It also outlines how the estimator relates to staged FEP and describes the simulation and sampling setup.
VII. DISCUSSION
- On a periodic soft-sphere solvation prototype, LFEP and LBAR reduced variance relative to their respective baselines; the direct approach was data-efficient in training and evaluation.The authors suggest it may be preferable when one state is a small perturbation of the other, and propose that mapping could be combined with staging.Source: §VII · pp. 7–8
- The proposed permutation-equivariant, periodic architecture addresses symmetries that the authors consider important for scaling learned free-energy estimation to atomistic system sizes.They note that a related gas-phase protein model lacked permutation equivariance, and identify comparisons with Boltzmann Generators and related estimators as future work.Source: §VII · p. 8
Appendix A: Alternate derivation of TFEP and interpretation
- The appendix interprets unidirectional TFEP as an FEP estimate between the mapped distribution A′ and state B, with convergence improved when A′ overlaps B.The derivation treats TFEP as a two-stage decomposition in which one free-energy contribution is zero.Source: Appendix A · p. 8
Appendix B: System
- For the solvation setup, 20 simulations per state yielded 1×10^6 samples, while MBAR used 13 intermediate states and 1×10^5 samples per state.The system used a cubic periodic box and a stationary solute; the MBAR calculation included 15 states in total.Source: Appendix B · pp. 8–9
Appendix C: Model implementation details
The model uses circular splines and a transformer-based network to implement mappings that satisfy the stated boundary and permutation symmetries.
Appendix C: Model implementation details
- Circular splines implement G as strictly increasing rational-quadratic segments, with endpoints set to satisfy the stated boundary conditions.Increasing the number of segments increases flexibility; sufficiently many can approximate any strictly increasing function meeting those conditions.Source: Appendix C · p. 9
- Transformer blocks implement C, using shared per-particle transformations and self-attention to preserve permutation equivariance and produce a permutation-invariant transformed distribution.Layer normalization is applied to the inputs of attention and MLP layers and to the output.Source: Appendix C · p. 9
- Table I summarizes the model and training hyperparameters.Source: Table I · p. 9
Appendix D: Results for LFEP
The LFEP mapping reduced variance across the evaluated training steps, but agreement with MBAR depended on when the mapping was evaluated. The KL minimum did not identify the lowest-bias mapping, consistent with the authors’ proposed zero-forcing explanation.
Appendix D: Results for LFEP
- The KL estimate reached a pronounced minimum while test and training losses were still decreasing, as the ΔF estimate drifted lower.Source: Appendix D · p. 9
- The KL minimum was not the lowest-bias stopping point; the authors suggest LLFEP’s zero-forcing property may explain this mismatch because it does not encourage mass coverage.They note that inadequate mass coverage is known to negatively affect importance-sampling estimators.Source: Appendix D · p. 9
- 2×10^3 training steps was the only evaluated mapping agreeing with MBAR; all three reduced variance, while later mappings became increasingly biased.Figure 6 compares the ΔF running averages for FEP, MBAR, and LFEP.Source: Appendix D · p. 9, Figure 6 · p. 10