Source-linked AI summary
Transport map accelerated Markov chain Monte Carlo
Matthew Parno, Youssef Marzouk
TL;DR
MCMC can struggle to explore complex target distributions efficiently, motivating proposals that capture their global structure rather than relying only on local or derivative information. The paper combines adaptively constructed transport maps with Metropolis-Hastings, using convex sample-based optimization and sequential updates. The resulting adaptive scheme preserves exact-target ergodicity and can be substantially more efficient on strongly non-Gaussian inference problems, although map approximation and early exploration constrain performance.
Problem
MCMC proposals may mix poorly when target distributions have nonlinear dependencies or non-Gaussian structure, while some competing methods require local derivative information.
Method
The paper adaptively constructs lower triangular approximate transport maps from MCMC samples and integrates them with Metropolis-Hastings to transform target sampling toward a simple reference distribution.
Results
The adaptive transport-map approach is more efficient on strongly non-Gaussian problems, with gains becoming more significant on the ODE and DAE examples.
Takeaways & Limitations
Approximate transport maps can support exact-target adaptive MCMC while exploiting global nonlinear and non-Gaussian posterior structure.
Takeaways & Limitations
The map is approximate because the expectation is estimated and the optimization domain and map parameterization are restricted; poor early exploration can also hinder adaptation.
Abstract
from arXiv · showhide
We introduce a new framework for efficient sampling from complex probability distributions, using a combination of optimal transport maps and the Metropolis-Hastings rule. The core idea is to use continuous transportation to transform typical Metropolis proposal mechanisms (e.g., random walks, Langevin methods) into non-Gaussian proposal distributions that can more effectively explore the target density. Our approach adaptively constructs a lower triangular transport map-an approximation of the Knothe-Rosenblatt rearrangement-using information from previous MCMC states, via the solution of an optimization problem. This optimization problem is convex regardless of the form of the target distribution. It is solved efficiently using a Newton method that requires no gradient information from the target probability distribution; the target distribution is instead represented via samples. Sequential updates enable efficient and parallelizable adaptation of the map even for large numbers of samples. We show that this approach uses inexact or truncated maps to produce an adaptive MCMC algorithm that is ergodic for the exact target distribution. Numerical demonstrations on a range of parameter inference problems show order-of-magnitude speedups over standard MCMC techniques, measured by the number of effectively independent samples produced per target density evaluation and per unit of wallclock time.
1. Introduction.
The paper proposes adaptive transport maps as a new way to design Metropolis-Hastings algorithms, transforming complex targets toward simple reference distributions. It motivates this approach by the limitations of local or derivative-dependent proposals and frames transport as a global alternative for improving MCMC exploration.
- 1. Introduction.: Transport maps provide monotone, typically nonlinear transformations between a target distribution and a simple reference distribution, acting like a preconditioner for MCMC.The paper uses these maps to make complex target distributions easier to sample.
- 1. Introduction.: Effective MCMC proposals seek highly independent successive states because better mixing reduces the variance of expectation estimates.Proposals may approximate the target locally, as in random-walk or Langevin methods, or globally, as in independence samplers.
- 1. Introduction.: Adaptive Metropolis methods update proposal covariance using empirical target information, while other methods use local gradients or higher-order derivatives.The paper contrasts covariance adaptation with HMC and stochastic Newton methods, which require derivative information.
- 1. Introduction.: The proposed approach adaptively constructs approximate transport maps from MCMC samples and integrates them with MCMC to retain exact sampling from the target.The maps represent the target globally as the pullback of a Gaussian reference measure.
- 1. Introduction.: The map-construction optimization is convex regardless of the target distribution, separable across dimensions, parallelizable, and does not require target-density derivatives.The method instead uses MCMC samples to construct the map and supports efficient sequential updates.
2. Construction of transport maps.
The paper constructs approximate lower triangular transport maps from target samples to a simple reference distribution. The resulting optimization is convex and separable, while approximation arises from sample averaging and finite-dimensional map restrictions.
- Optimal transportation: A transport map deterministically pushes the target measure to a reference measure, with the Knothe-Rosenblatt rearrangement providing a uniquely defined lower triangular map under absolute continuity.Its Jacobian is lower triangular with positive diagonal entries.
- Optimal transportation: Lower triangular maps are imposed directly rather than constructed through the weighted transport cost, avoiding numerical underflow as dimension increases.This differs from standard optimal transportation while preserving the desired map structure.
- Optimal transportation: Approximate maps can capture key target structure and enable more efficient MCMC proposals when exact maps are computationally challenging.An exact map would allow reference samples to be transformed directly into target samples.
- Optimization objective: The map-induced density is compared with the target density through KL divergence, using a sample-average approximation to obtain a deterministic optimization problem.The constructed map is approximate because both the expectation and feasible map class are restricted.
- Constraints: A positive lower bound on each diagonal map derivative makes the Jacobian positive definite and removes the absolute value from determinant terms.The paper typically chooses λmin = 10^-8 without tuning it in practice.
- Convexity and separability: The optimization separates into n individual convex problems because both the objective and constraints involve one map component at a time.Each component is parameterized with a multivariate polynomial expansion.
3. Map-based MCMC proposals.
Transport maps modify Metropolis-Hastings by applying proposals in a reference space and pulling them back to the target space. The resulting proposal can use local or independence mechanisms while requiring target derivatives only when the reference proposal does.
- Map-based MCMC proposals: A fixed transport map transforms either the target distribution or the proposal mechanism within Metropolis-Hastings.The adaptive construction of such maps is developed separately.
- Map-based MCMC proposals: The map encodes target information so reference-space proposals can make larger moves while retaining the Metropolis-Hastings accept/reject correction.Proposal choice controls dependence between successive chain states through acceptance rate and step size.
- Proposal process: The reference proposal may be local or independent because it is applied to the pushforward distribution, which is closer to a simple reference distribution.The pushforward is not necessarily Gaussian, but the map is chosen to make it closer to standard Gaussian.
- Proposal process: The algorithm maps the current target state to reference space, samples a reference proposal, and applies the inverse map to obtain the target-space proposal.These are the three proposal-sampling steps in Algorithm 1.
- Implementation: Lower triangularity reduces inverse-map evaluation to n one-dimensional nonlinear solves.The implementation uses a robust bisection solver based on Sturm sequences for polynomial maps.
- Derivative requirements: Target-density derivatives are required if and only if the chosen reference proposal explicitly requires them.The method does not require a particular target form and can handle targets without available derivative information.
4. Adaptive transport map MCMC.
The algorithm adapts a transport map from prior MCMC states, using regularization to stabilize early updates and efficient Newton iterations for sequential recalculation. Reference-space proposals, including local, independence, Langevin, and delayed-rejection schemes, induce corresponding target-space proposals.
- Adaptive algorithm overview: Previous MCMC states update a nonlinear transport map every KU iterations, replacing covariance adaptation with more general non-Gaussian proposals.The map is initialized with eT0 and updated by solving the sample-based optimization problem.
- Adaptive algorithm overview: Regularization prevents early map collapse when few samples or poor mixing make the Monte Carlo objective unreliable.The penalty encourages the map to remain near the identity, especially when K is small.
- Adaptive algorithm overview: A quadratic penalty centered on identity-map coefficients is used in practice, with kR mostly set to 10^-4 in numerical examples.As the sample count grows, the Monte Carlo term overwhelms regularization and kR becomes less important.
- Adaptive algorithm overview: Sequential map updates are computationally efficient because nearby objectives make previous coefficients good Newton initializations, sometimes requiring only one or two iterations.The reported timing results indicate that map-proposal benefits outweigh sequential update costs even for long chains.
- Choice of reference proposal: The reference proposal can be random-walk, independence, Langevin, or delayed-rejection, while the map transforms it into a target-space proposal.Delayed rejection combines a broad first-stage proposal with a more conservative second stage.
- Choice of reference proposal: Figure 2 compares local and independence proposals in reference and target spaces, showing target densities as black contours and proposals as colored contours.The target-space proposal is qθ(θ′|r) = qr(r′|r)|D eT(θ′)|.
5. Convergence analysis.
The convergence analysis identifies conditions under which adaptive map MCMC remains ergodic for the exact target distribution. The proof combines bounded map-induced proposals with diminishing adaptation and drift-based control, while nonlinear maps require derivative restrictions.
- Failure mode: Nonlinear polynomial maps of degree greater than one can destroy super-exponential tail behavior even for a Gaussian target.The resulting loss of geometric ergodicity is linked to unbounded map derivatives.
- Conditions for convergence: Bi-Lipschitz restrictions on the map are needed to control derivatives and restore convergence guarantees.The admissible parameter space is restricted so every map satisfies the stated condition.
- Proof strategy: The proof establishes bounded target-space proposals through upper and lower bounds inherited from the map conditions.These bounds support the subsequent minorization and drift arguments.
- Proof strategy: The adaptive chain satisfies diminishing adaptation when map parameters are restricted to a compact subset of the feasible parameter space.This controls changes in the transition kernels across updates.
- Proof strategy: Minorization and drift conditions establish simultaneous strongly aperiodic geometric ergodicity, which supplies containment for the adaptive chain.The convergence proof combines these properties with diminishing adaptation.
- Conclusion: Under finite, continuous, super-exponentially light targets and constrained maps, Algorithm 1 is ergodic for the target distribution.The theorem requires map parameters to lie in a compact set where the bi-Lipschitz condition holds.
6. Numerical examples.
Numerical examples show that adaptive transport-map proposals improve sampling across challenging posterior geometries, while performance depends on initialization, map expressiveness, and proposal type.
- Experimental setup: The study compares transport-map MCMC variants with DRAM, sMMALA, AMALA, and NUTS using minimum effective sample size across parameter dimensions.Performance is evaluated through integrated autocorrelation time and effective sample size across independent chains.
- Biochemical oxygen demand model: Cubic Hermite transport maps capture changing posterior correlations and non-Gaussian structure that challenge standard samplers.For the BOD problem, global independence proposals can jump across the parameter space after the target is transformed toward a more Gaussian distribution.
- Biochemical oxygen demand model: TM+DRG produced much shorter integrated autocorrelation times in the BOD example by using global independence proposals after adaptation.The map captures global non-Gaussian structure, enabling efficient moves across the entire parameter space.
- Biochemical oxygen demand model: TM+LA performed poorly because weak early exploration hindered adaptation from the initial identity map.The authors therefore excluded TM+LA from the remaining test problems.
- Predator-prey system: Longer predator-prey chains favored transport-map approaches, with map-accelerated chains showing significantly better mixing and improvements of at least two orders of magnitude over baseline schemes.The authors attribute the longer-chain advantage to reduced relative burn-in and additional time for adaptive proposals to characterize the posterior.
- Maple sap exudation: For the maple posterior, nonlinear dependencies yielded order-of-magnitude gains, but TM+DRL outperformed TM+DRG because cubic maps could not fully characterize the target.Local proposals were less sensitive to deficiencies in the map, whereas global independence proposals were less effective.
- Adaptation limitations: Challenging targets can cause proposal collapse when small initial steps poorly explore the space and regularization is insufficient.The resulting map may place excessive probability mass on the limited region visited by the initial chain.
7. Conclusions.
The approach gains efficiency by capturing nonlinear and non-Gaussian posterior structure, while map construction adds computational cost and leaves substantial parallelism unused. Its representation is flexible beyond polynomial bases, with adaptive basis selection identified as important for higher-dimensional problems.
- Capturing nonlinear dependencies and non-Gaussian posterior structure primarily drives efficiency, especially when global independence proposals can exploit that information.The paper reports gains in effective sample size per posterior evaluation and per unit of wallclock time.
- Sequential transport-map updates add computational cost that may matter in simple problems.
- The current implementation does not exploit parallelism across parameter dimensions and samples during map construction.The optimization can be made embarrassingly parallel over parameter dimensions, with additional parallelism over samples.
- Polynomial representation is not essential because the map-coefficient optimization remains unchanged for representations linear in their coefficients.The paper mentions radial basis functions as an alternative basis.
- Adaptive basis selection is especially relevant for extending the approach to higher-dimensional problems, where more parsimonious bases will be needed.
Appendix A. ESS Calculation Details.
The appendix defines maximum integrated autocorrelation time and effective sample size across multiple chains and dimensions. Normalized ESS metrics include burn-in costs through their evaluation or wallclock-time denominators.
- Maximum integrated autocorrelation time takes the median across chains and the maximum across parameter dimensions.Dimension-specific autocorrelation times are computed independently for each chain and dimension using a Fourier-transform method.
- For dimension j on chain i, ESS_i,j equals K^2 divided by 2τ_i,j, where K is the number of post-burn-in samples.
- The implementation and source code are associated with MUQ, using a specified repository commit.
- Normalized ESS per function evaluation and per second uses all evaluations or computational time in the denominator, including burn-in costs.
Appendix B. Proof of ergodicity.
The appendix supplies technical analysis supporting the paper’s convergence results, including detailed proofs of two lemmas and extensions used for a later theorem.
- The appendix provides technical analysis associated with the convergence properties of the map-accelerated MCMC algorithm.
- It gives detailed proofs of Lemmas 5.3 and 5.4, while the remaining ingredients for Theorem 5.7 extend a proof from prior work.
B.1. Bounded target proposal.
The bounded-target-proposal analysis bounds the map-induced proposal using Gaussian densities under regularity assumptions on the target, reference proposal, and transport map. These bounds establish Lemma 5.3.
- The analysis assumes a finite, continuous, super-exponentially light target density and a Gaussian random-walk reference proposal with bounded drift and fixed covariance.
- Two zero-mean Gaussian densities and finite positive scalars are constructed to support bounds on the proposal analysis.
- The bi-Lipschitz condition on the transport map is used to derive an upper bound for the target-space proposal density.
- The corresponding lower bound follows through an analogous argument, and together the bounds imply Lemma 5.3.
B.2. SSAGE.
This section establishes SSAGE by combining a minorization result with a drift condition for the adaptive Metropolis-Hastings kernel. The proof uses guaranteed-acceptance and possible-rejection regions, together with tail assumptions on the target density.
- SSAGE: SSAGE reduces to proving the minorization component and the drift condition after the preceding lemmas establish the required acceptance properties.The minorization component is provided by Lemma 5.5, leaving Lemma 5.6 as the remaining requirement.
- SSAGE: The proposal has a nonzero probability of acceptance, established by constructing a measurable accepted subset within a small neighborhood of the current state.The argument uses positive-variance, zero-mean proposal components and a small ball around x.
- SSAGE: For sufficiently large states, super-exponential tails and curvature ensure that a cone of proposals closer to the origin is accepted.The cone is contained in the guaranteed-acceptance region because the target density is larger there than at an already accepted point.
- SSAGE: The drift proof uses V(x) = cV π^-α(x), with α ∈ (0, 1), and decomposes the Metropolis-Hastings transition over rejection and acceptance regions.Acceptance rates in the possible-rejection region lie in [0, 1), enabling the subsequent bounds.
- SSAGE: Bounds derived from the two regions satisfy the limiting inequalities required for the drift condition, completing Lemma 5.6.The argument combines Lemmas B.1 and B.2 with the bounds corresponding to the acceptance and rejection regions.
B.3. Diminishing adaptation.
This section proves diminishing adaptation by showing that transition kernels become arbitrarily close as the sample size used for map construction increases. The argument relies on continuity of the map and convergence properties of the finite-sample optimization problem.
- Diminishing adaptation: The map coefficients minimize a K-sample Monte Carlo approximation of the KL divergence, without requiring convergence to the true KL minimizer.The proof explicitly uses the finite-sample objective and does not assume that it converges to the exact population minimizer.
- Diminishing adaptation: Diminishing adaptation requires the difference between transition kernels at sample sizes K and K + KU to vanish as K →∞.When no adaptation occurs, consecutive maps are identical, so only changes across adaptation steps require control.
- Diminishing adaptation: Continuity of the map parameters and bounded proposal densities reduce diminishing adaptation to convergence of the optimized map coefficients.The maps are linear in the coefficients, making the transition probability continuous for fixed x and measurable A.
- Diminishing adaptation: The resulting bounds establish the desired diminishing-adaptation condition, while the proof also uses finite expectations ensured by map regularity and Gaussian proposal bounds.The map is bi-Lipschitz, the proposal density is Gaussian-bounded, and the map is linear for large ∥θ∥.
- Diminishing adaptation: Convexity and uniqueness of each objective’s global minimizer cause the relevant coefficient sets to collapse around γ(K) as δ2 →0.The shrinking sets control the maximum distance between candidate coefficient vectors.
Appendix C. Maple exudation model details.
This appendix specifies the maple sap exudation forward model used in the parameter-inference example. It includes the model’s physical variables, heat equation, piecewise conductivity, numerical solver, and parameter transformations.
- Numerical solution: The system is solved with MUQ linked to SUNDIALS, using steady-state initial conditions for selected state variables and U(0) = 0.The initial conditions are specified for sgi, siw, and r(t).
- Thermal model: The temperature field is modeled as quasi-steady and satisfies a one-dimensional heat equation over the time-dependent computational domain.The equation uses thermal conductivity κ(x) and spans from siw(t) to Rf + 2Rv.
- Thermal model: Thermal conductivity is piecewise defined using water and air conductivities across a region determined by r(t), producing a piecewise linear temperature field.At each time, the heat equation can be solved analytically.
- Parameterization: Inference parameters θ are transformed into the physical model parameters using Table 4, with default parameters and remaining constants specified separately.Variables with overbars are default parameters taken from the original model reference.