Flow Perturbation++: Multi-Step Unbiased Jacobian Estimation for High-Dimensional Boltzmann Sampling

TL;DR

Flow Perturbation++ introduces a multi-step, variance-reduced Jacobian estimator for high-dimensional Boltzmann sampling, significantly improving efficiency and accuracy in complex systems like 1000D Gaussian mixtures and proteins.

Contribution

It extends Flow Perturbation with a multi-step approach that reduces variance while maintaining unbiased Jacobian estimation in high-dimensional sampling.

Findings

Achieves lower estimator variance compared to previous methods.

Improves equilibrium sampling accuracy in high-dimensional models.

Demonstrates effectiveness on protein and Gaussian mixture benchmarks.

Abstract

The scalability of continuous normalizing flows (CNFs) for unbiased Boltzmann sampling remains limited in high-dimensional systems due to the cost of Jacobian-determinant evaluation, which requires $D$ backpropagation passes through the flow layers. Existing stochastic Jacobian estimators such as the Hutchinson trace estimator reduce computation but introduce bias, while the recently proposed Flow Perturbation method is unbiased yet suffers from high variance. We present \textbf{Flow Perturbation++}, a variance-reduced extension of Flow Perturbation that discretizes the probability-flow ODE and performs unbiased stepwise Jacobian estimation at each integration step. This multi-step construction retains the unbiasedness of Flow Perturbation while achieves substantially lower estimator variance. Integrated into a Sequential Monte Carlo framework, Flow Perturbation++ achieves significantly improved equilibrium sampling on a 1000D Gaussian Mixture Model and the all-atom Chignolin protein compared with Hutchinson-based and single-step Flow Perturbation baselines.