Flow Perturbation to Accelerate Unbiased Sampling of Boltzmann distribution

TL;DR

This paper introduces a flow perturbation method that significantly accelerates unbiased sampling of the Boltzmann distribution in high-dimensional systems, enabling detailed molecular sampling previously infeasible with generative models.

Contribution

The authors propose a novel flow perturbation technique that reduces computational costs and achieves unbiased Boltzmann sampling, demonstrated on large molecular systems.

Findings

Achieved orders of magnitude speedup over traditional Jacobian calculations.

Successfully sampled the atomic coordinates of the large protein Chignolin.

Demonstrated unbiased sampling of high-dimensional molecular systems.

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.