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

TL;DR

This work presents 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, which retains the unbiasedness of Flow Perturbation while achieves substantially lower estimator variance.

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.

Figures (6)

• Figure 1: (a) Single-step Flow Perturbation (FP) computes the full flow Jacobian in one shot, often resulting in high variance. (b) Flow Perturbation++ (FP++) breaks the flow into $T$ steps, applying unbiased estimation sequentially from latent $\mathbf{z} = \mathbf{x}_{T+1}$ to the final sample $\mathbf{x} = \mathbf{x}_1$. The product of per-step Jacobians gives a low-variance, numerically stable approximation of the full determinant.
• Figure 2: Representative CNL025 configurations at 300 K: (a) partially $\alpha$-helical; (b) folded $\beta$-hairpin.
• Figure 3: Final SMC sample distribution along the reaction coordinate for the 10D GMM. The solid blue curve shows the target distribution obtained via direct sampling.
• Figure 4: SMC results for the 1000-dimensional GMM. (a) Probability distribution of final SMC samples along the reaction coordinate. (b) Monte Carlo acceptance rate as a function of $\beta$. (c) Average sample energy as a function of $\beta$. (d) Final energy distribution. The target energy and distribution (solid line) are obtained by direct sampling from the target GMM. All results are derived from a single SMC run.
• Figure 5: SMC results for the Chignolin mutant. (a–f) Sample distributions projected onto the reaction-coordinate (RC) plane: (a) Target distribution derived from the training data; (b) Samples generated directly by the base flow model; (c) Samples produced by SMC using the single-step FP estimator; (d) Samples produced by SMC using the FP++ estimator; (e) Samples produced by SMC with the Hutchinson estimator (10 Rademacher vectors); (f) Samples produced by SMC with the Hutchinson estimator (10 Gaussian vectors). (g) Monte Carlo acceptance rates as functions of $\beta$; (h) Average energy $\langle E \rangle$ of intermediate SMC samples; (i) Energy distributions at the end of SMC compared to the target distribution and the base flow (“init”). All SMC results are derived from a single run.