Flow Perturbation to Accelerate Unbiased Sampling of Boltzmann distribution

TL;DR

The flow perturbation method is introduced, which incorporates optimized stochastic perturbations into the flow and achieves unbiased sampling of the Boltzmann distribution with orders of magnitude speedup compared to both brute force Jacobian calculations and the Hutchinson estimator.

Abstract

Flow-based generative models have been employed for sampling the Boltzmann distribution, but their application to high-dimensional systems is hindered by the significant computational cost of obtaining the Jacobian of the flow. To overcome this challenge, we introduce the flow perturbation method, which incorporates optimized stochastic perturbations into the flow. By reweighting trajectories generated by the perturbed flow, our method achieves unbiased sampling of the Boltzmann distribution with orders of magnitude speedup compared to both brute force Jacobian calculations and the Hutchinson estimator. Notably, it accurately sampled the Chignolin protein with all atomic Cartesian coordinates explicitly represented, which, to our best knowledge, is the largest molecule ever Boltzmann sampled in such detail using generative models.

Figures (14)

• Figure 1: Reweighting trajectories generated by a flow model. The prior distribution, shown in the middle, is Gaussian, while the target distribution features two peaks of different heights. The original trajectories, shown on the left, fail to sample the target distribution correctly, with roughly equal numbers of trajectories reaching each peak despite their differing heights. On the right, the trajectories are reweighted according to $e^{-W}$, resulting in trajectories that arrive at the lower probability regions being given less weight. This reweighting ensures unbiased sampling of the target distribution.
• Figure 1: Neural Network structure of $F_{\theta}$ for GMM
• Figure 2: An illustration of the flow perturbation method. The left panel shows the forward and corresponding backward trajectories generated by the original flow model, while the right panel shows the trajectories generated by the perturbed flow model.
• Figure 2: Neural Network structure of $\epsilon_{\theta}$ for Chignolin
• Figure 3: The representative configurations of Chignolin at 300K: (a) beta-hairpin and (b) a configuration with partial alpha-helical segment.