Source-linked AI summary
Adaptive Monte Carlo augmented with normalizing flows
Marylou Gabrié, Grant M. Rotskoff, Eric Vanden-Eijnden
TL;DR
High-dimensional multimodal sampling is difficult for local MCMC, especially when metastable basins are separated by large free-energy barriers and their relative weights are unknown. The paper develops adaptive MCMC that combines local and normalizing-flow proposals, trains the flow on generated samples, and theoretically analyzes convergence. Experiments show effective cross-basin sampling and dramatic acceleration, while performance depends on prior information about basin locations and suitable flow design.
Problem
Local MCMC struggles with high-dimensional multimodal targets, where estimating relative metastable-state weights can require immense computation or carefully designed importance sampling.
Method
The method blends local MCMC moves with nonlocal normalizing-flow proposals and adapts the flow during sampling using generated data.
Results
Experiments show dramatically accelerated sampling across metastable states and efficient estimation of their relative statistical weights without constructing specific pathways between basins.
Takeaways & Limitations
Normalizing-flow augmentation can accelerate sampling across large free-energy barriers when the rough locations of relevant metastable modes are known beforehand.
Abstract
from arXiv · showhide
Many problems in the physical sciences, machine learning, and statistical inference necessitate sampling from a high-dimensional, multi-modal probability distribution. Markov Chain Monte Carlo (MCMC) algorithms, the ubiquitous tool for this task, typically rely on random local updates to propagate configurations of a given system in a way that ensures that generated configurations will be distributed according to a target probability distribution asymptotically. In high-dimensional settings with multiple relevant metastable basins, local approaches require either immense computational effort or intricately designed importance sampling strategies to capture information about, for example, the relative populations of such basins. Here we analyze an adaptive MCMC which augments MCMC sampling with nonlocal transition kernels parameterized with generative models known as normalizing flows. We focus on a setting where there is no preexisting data, as is commonly the case for problems in which MCMC is used. Our method uses: (i) a MCMC strategy that blends local moves obtained from any standard transition kernel with those from a generative model to accelerate the sampling and (ii) the data generated this way to adapt the generative model and improve its efficacy in the MCMC algorithm. We provide a theoretical analysis of the convergence properties of this algorithm, and investigate numerically its efficiency, in particular in terms of its propensity to equilibrate fast between metastable modes whose rough location is known \textit{a~priori} but respective probability weight is not. We show that our algorithm can sample effectively across large free energy barriers, providing dramatic accelerations relative to traditional MCMC algorithms.
I. INTRODUCTION
High-dimensional multimodal distributions challenge local MCMC because metastable basins are separated by low-probability regions. The paper proposes adaptive MCMC with normalizing-flow proposals, analyzes convergence, and reports accelerated sampling when basin locations are known approximately but weights are not.
- Metastability makes local MCMC struggle to transition between high-probability basins separated by low-probability regions.
- The paper proposes adapting a normalizing flow during MCMC sampling, without requiring an extensive preexisting dataset.
- The algorithm combines local MCMC moves with nonlocal flow proposals and trains the flow using generated samples.
- Theoretical analysis treats simultaneous training and sampling as a nonlinear MCMC scheme with asymptotic exponential convergence in continuous time.
- Numerical experiments show dramatic acceleration and efficient estimation of relative statistical weights across metastable states without constructing pathways between basins.
- The approach requires representative configurations for modes of interest beforehand, while generative proposals alone are statistically unlikely to discover unknown metastable states.
III. MCMC SAMPLING WITH NORMALIZING FLOWS
Normalizing flows provide invertible proposals whose densities can be evaluated for Metropolis-Hastings correction. Their effectiveness depends on overlap with the target, motivating training from limited data and careful map design.
- A normalizing flow transports samples from a tractable base distribution to a target distribution through an invertible map.
- An approximate flow can still produce exact target sampling when its proposals are corrected with a Metropolis-Hastings accept-reject step.
- Flow proposals may mix faster than local moves because they can directly generate independent configurations from the push-forward distribution.
- Poor overlap between the push-forward and target distributions can make acceptance rates very low, especially in high-dimensional spaces.
- Training solely from base-distribution samples can be difficult because relevant target regions may initially have exponentially small weight.
- The adaptive procedure trains the map using generated samples and updates it through an optimization step based on the estimated forward KL divergence.
IV. ADAPTIVE MCMC: CONCURRENT SAMPLING AND TRAINING
The adaptive scheme concurrently samples with local and flow-based kernels while using the resulting data to optimize the flow. Combining both kernels preserves detailed balance and improves robustness during imperfect training.
- Sampling and flow optimization proceed concurrently, so newly acquired data further improve the nonlocal proposal distribution.
- The algorithm combines a local kernel with a normalizing-flow kernel, and any detailed-balance MCMC method can serve as the local component.
- The compounded transition kernel satisfies detailed balance because its local and flow-based kernels individually satisfy detailed balance.
- Alternating local and flow moves improves robustness when the flow is initially untrained or cannot represent every target feature.
- The compounded kernel has a leading eigenvalue bounded by the product of the individual leading eigenvalues, indicating faster convergence under the stated spectral-gap assumptions.
B. Adaptation
The method adapts a normalizing-flow proposal using configurations generated during MCMC, while combining nonlocal resampling with local dynamics. Its analysis shows asymptotic convergence with exponential decay controlled by the resampling component under stated assumptions.
- Adaptation: The flow is trained on configurations sampled from the current chain by minimizing an estimated forward KL divergence to the flow distribution.The estimator uses the chain samples and stochastic-gradient updates to optimize the normalizing-flow parameters.
- Initialization: Chains begin from configurations located in the target’s different modes, although these initial configurations need not be drawn from the target distribution.The procedure assumes prior information about metastable-state locations and should not be expected to find basins distinct from initialization.
- Adaptation: Initially, local MCMC drives sampling because the flow starts as the identity; as training improves, accepted nonlocal moves reduce autocorrelation and reallocate mass according to basin weights.This progression is illustrated for Gaussian-mixture and random-field examples.
- Convergence: In the continuous-time analysis, α controls the balance between Langevin dynamics and resampling, with α = 0 corresponding to Langevin dynamics alone.Langevin-only convergence is asymptotically valid but generally exponentially slow.
- Convergence: After an approximately constant initial period, the Pearson χ2-divergence decays exponentially at rate 2α > 0 under the instantaneous-distribution assumption.The constant-period condition is α(t − t0) ≤ log G^-1, after which the resampling step controls the exponential rate; concurrent Langevin steps can only improve it.
E. Scalability: Model-Informed Base Distributions and Maps, Mixtures, etc
High-dimensional flow learning is eased by incorporating physics-informed base measures, map architectures, and mixtures tailored to metastable modes. In the Allen-Cahn example, these choices enable accurate proposals, rapid mixing, and free-energy estimation across conditions.
- Model-informed distributions and maps: Physics-informed base measures and map classes ease learning when target distributions have anisotropic correlations or system-specific structure.The paper uses covariance-matched Gaussian fields, physics-informed distributions, and factorized maps tailored to interacting particles.
- Mixtures: Mixtures of separately pretrained maps represent modes with different fine structures or statistical weights by adapting their mixture weights during sampling.Each component is trained from local-MCMC data initialized in its corresponding mode before adaptive sampling begins.
- Allen-Cahn model: Around 60% of informed-flow proposals are accepted, whereas the uninformed flow fails to generate samples with appreciable target statistical weight.The informed base measure captures high-frequency target features that the uncoupled base measure misses, causing high rejection rates.
- Allen-Cahn model: At N = 100, preadapting the push-forward covariance is highly beneficial when no empirical whitening transform can be computed from preexisting samples.The proposed informed base measure supplies this preadaptation using knowledge of the target Hamiltonian.
- Allen-Cahn model: The learned Allen-Cahn map is spatially local, transports the mean mode toward ±1, and leaves higher-frequency modes approximately invariant.This structure helps explain how the map captures the system’s large-scale phase organization without altering higher-frequency components.
- Free-energy estimation: Importance sampling with the learned flow accurately recovers free-energy differences between metastable basins under changed biasing fields and temperatures.The method generalizes when the modified parameter changes basin populations while only mildly affecting local field structure.
B. Detecting Phase Transitions in Interacting Particle Systems
Interacting particle systems exhibit exponentially slow brute-force transitions between metastable gas and liquid phases. The adaptive mixture proposal learns their relative weights and detects the first-order transition despite this trapping.
- Particle system and phase diagram: In the thermodynamic limit, the model has a first-order transition between a uniform gas phase and a clustered liquid phase.The phase diagram is estimated using a mean-field approximation.
- Particle system and phase diagram: Brute-force transition times diverge exponentially with particle number because free-energy barriers trap the system in either metastable phase.The reported scalings are t_l→g ≍ exp(NβF_l→g) and t_g→l ≍ exp(NβF_g→l).
- Adaptive simulations: The nonlocal proposal uses a mixture: uniform particle positions for the homogeneous phase and a trained map that transforms uniform configurations into droplet-like configurations.This parameterization accommodates phases with different structures and statistical weights.
- Adaptive simulations: The Metropolis-Hastings ratio combines energy and entropy contributions, requiring the normalizing-flow density to remain sufficiently accurate for acceptance.The exponential factor must be of order 1 in N to obtain significant acceptance, even for moves toward the preferred phase.
- Implementation and results: Starting from p(0) = 1/2, the adaptive procedure learns p values converging to 0 or 1 and correctly reweights homogeneous and droplet configurations.It determines the more likely phase even where brute-force molecular dynamics observes no transitions.
- Implementation and results: The particle demonstration assumes no short-range repulsion; hardcore systems require more complex mixtures or flow maps and remain future work.The product-density mixture used here is invalid when particles exhibit hardcore repulsions.
VI. CONCLUSIONS
The paper positions adaptive normalizing-flow MCMC as a useful nonlocal sampling framework, while emphasizing limitations in mode discovery, scaling, and architectural generality.
- Reverse-KL flow training can suffer mode collapse by concentrating on one mode and missing distant modes.The failure occurs when the map does not explore regions far from the base distribution's core.
- The proposed method extends Markov Score Climbing by alternating forward-KL adaptation with local sampling and nonlocal resampling.The local component distinguishes the approach from latent-space local proposals.
- Nonlocal transport and normalizing flows can enhance mixing even without an initially representative dataset.The conclusion still assumes some prior information about metastable states.
- The method does not enable discovery of unknown metastable modes, so prior information about their locations remains important.This is the paper's main stated scope boundary.
- Future scaling depends on efficient learning and map properties such as locality, transferability, and incorporation of target symmetries.The authors identify symmetry-aware architectures as an active research direction.
Appendix A: Comparing the Various Methods of Training and Sampling
The appendix contrasts flow training with target data against self-training from a known energy, then explains why combining local MCMC with flow resampling is more robust when the map is initially poor.
- Training with target data: With target samples but no analytical density information, direct-KL training fits the flow using empirical target data.Sampling then pushes base-distribution points through the trained map, but generated samples cannot be independently cross-validated.
- Self-training from target energy: With a known target energy but no target samples, reverse-KL self-training optimizes the flow using samples generated from the base distribution.The base distribution can provide arbitrarily many training points, but poor initial maps can make useful target regions exponentially rare.
- Self-training from target energy: A poorly initialized map can make self-training difficult or impossible and yield flow proposals with very low Metropolis-Hastings acceptance.The problem arises because target-relevant samples initially receive exponentially small weight under the base distribution.
- Blending local and nonlocal sampling: Starting local MCMC walkers in approximately known high-probability regions supplies locally accurate samples for direct-KL flow training without their relative basin weights.The subsequent nonlocal moves equilibrate basin weights and provide better proposals for further map refinement.
- Blending local and nonlocal sampling: Keeping local MCMC remains beneficial because it captures target details that an imperfect or insufficiently expressive flow misses.Pure flow resampling is appropriate only if the learned map closely approximates the optimal transport map.
- Measure-theoretic setup: The framework uses invertible maps to transport an efficiently sampled base measure toward the target and requires mutual absolute continuity for Metropolis-Hastings use.The KL objective must also be finite initially for the optimization argument to apply.
- Application-specific base measures: For stochastic Allen-Cahn problems, the appendix uses a Gaussian Ornstein-Uhlenbeck bridge as the base measure and estimates training quantities from current MCMC samples.The fields are discretized on a grid and integrals replaced by Riemann sums in practice.
2. Transition Paths
For transition-path sampling, the method uses a bridge-based Gaussian base measure and combines flow resampling with local MCMC updates in an alternating procedure.
- Base measure: A Gaussian bridge connecting xA to xB can serve as the base measure for transition-path sampling.The bridge is built from a standard Brownian bridge in the relevant state dimension.
- Base measure: The bridge construction supports a Radon-Nikodym derivative when the drift satisfies the Novikov condition.This condition ensures the stated martingale property.
- Transition kernel: The flow-based transition kernel updates the walker density through a Chapman-Kolmogorov equation.The kernel is analyzed using the target density and the flow-generated density.
- Transition kernel: If the generated density equals the target density, the resampling kernel reaches equilibrium in one step from any initial density.In this ideal case, the acceptance ratio correction is identically one.
- Algorithm: Algorithm 2 alternates partial flow resampling and local Langevin-style updates using the evolving map.The algorithm takes target energy, initial walkers, timestep, duration, and balance parameter as inputs.
2. Continuous limit
The continuous-limit analysis shows that flow resampling can accelerate convergence, while the guarantees depend on overlap, finite divergences, and initialization near metastable basins.
- Continuous evolution: The continuous-limit density evolution combines partial resampling with Langevin dynamics, with α controlling their relative timescale.The resulting equation is obtained by taking the timestep τ to zero.
- Convergence conditions: Convergence is guaranteed when the accumulated inverse overlap between the evolving flow density and target diverges, and is asymptotically exponential when overlap stays positive.The stated condition allows failure when the relevant lower bound decays faster than t^-1.
- Fixed-map rate: For a fixed flow, the convergence-rate bound is nontrivial only when its minimum density ratio ˆG is positive and can become poor when ˆG is exponentially small in dimension.The ideal value ˆG = 1 corresponds to an exact flow density equal to the target.
- Birth-death formulation: Under bounded Pearson χ2-divergences, the birth-death formulation admits convergence bounds for the sample density.The bounds require finite initial divergences, with a corresponding shifted-time version when finiteness begins at t0.
- Resampling interpretation: The adaptive flow-resampling term adds or removes mass according to the ratio between the current and target densities, then reinjects mass using the flow density.This mechanism stops when the relevant density functional reaches its target value.
- Scope boundary: Finite initial error depends on having samples near metastable basins, so the method is not generically suitable when those basins are unknown.This explicitly limits the scope of the convergence analysis and method.
Appendix E: Sampling unvisited metastable states
A normalizing flow trained on one metastable basin generally cannot discover an unvisited basin, making initialization with representatives from each basin necessary.
- A flow trained on one metastable basin is generally unsuitable for discovering unvisited metastable basins.
- The method therefore requires prior information about metastable-state locations to initialize chains across the modes of interest.
- The probability of sampling the opposite basin can be worse with a trained map than with the untrained identity map.
- At large system size, sampling the opposite basin is dominated by the fully aligned configuration and becomes a rare event at low temperature.
Appendix F: Normalizing flow architecture
The appendix describes RealNVP architectures and experiments showing how initialization, base measures, local updates, and flow expressiveness affect adaptive sampling.
- Normalizing flow architecture: RealNVP maps use invertible affine coupling layers whose inverse and Jacobian determinant can be computed efficiently.
- Normalizing flow architecture: In the Gaussian-mixture experiments, six coupling-layer pairs and three-layer, 100-unit multilayer perceptrons were used for scaling and translation networks.
- Two-dimensional wiggle: A two-dimensional wiggle experiment compares local MALA, flow-based NF, and alternating NF+MALA kernels when flow expressiveness is limited.
- Base measures and experiments: The uniformed prior fails to rebalance modes because almost no generated samples are accepted, and refinement makes fine-scale structure harder to represent.
- Two-dimensional wiggle: NF+MALA improves tail coverage and yields more credible effective-sample-size estimates near stationarity, while MALA more accurately samples the distribution initially.
- Additional experiments: The tilted Allen-Cahn experiment learns two domain-wall configurations, while transition-path experiments use Langevin-only training to mitigate premature learning of a higher-probability mode.
- Additional experiments: The learned map captures low-frequency path structure while higher-order terms restore high-frequency oscillations.
Appendix J: Interacting particle system experiments
Interacting-particle experiments contrast brute-force trapping with adaptive flow-assisted sampling across homogeneous and droplet phases, using gradient-based adaptation over temperatures.
- Mean-field formulation: In the mean-field limit, particle evolution follows a McKean–Vlasov equation that is a Wasserstein gradient flow of the free energy.
- Phase structure: The particle system exhibits gas-like homogeneous and liquid-like droplet phases separated by a first-order transition.
- Phase structure: Around the critical inverse temperature, the less-favored phase remains a local minimizer, indicating metastability at the first-order transition.
- Numerical procedure: The phase diagram is obtained numerically by iterating the density equation and assessing the stability of homogeneous and droplet states.
- Brute-force simulations: Brute-force simulations with 200 particles remain trapped in whichever homogeneous or droplet state they initially occupy.
- Adaptive simulations: Adaptive simulations alternate particle-dynamics steps with resampling and optimize the flow parameter concurrently, achieving average resampling acceptance rates between 35% and 55%.
Appendix K: Computing free energy differences
The appendix presents Metropolis-Hastings and importance-sampling estimators for free-energy differences, with effective sample size used to assess estimator quality.
- Metropolis-Hastings estimation: Free-energy differences between metastable configuration sets can be estimated from unbiased Metropolis-Hastings samples using a fast-mixing flow-based kernel.
- Estimator diagnostics: Effective sample size provides a quality monitor for both estimator types.
- Metropolis-Hastings estimation: The local kernel can compensate when the learned flow does not represent all target probability mass accurately.
- Importance sampling: Importance sampling uses unnormalized weights e^(-U*(x_i))/rho-hat(x_i), with the unknown partition function canceling in the free-energy difference estimator.
- Transfer across conditions: A map learned under one thermodynamic condition can be reused at distinct conditions because Metropolis-Hastings correction and importance weights account for target–proposal mismatch.