Transport RJMCMC via Normalizing Flows

• The paper introduces transport RJMCMC as a method that leverages deep normalizing flows to construct diffeomorphic maps for near-optimal transdimensional proposals.
• It transforms complex, dimension-changing moves into tractable operations in a common reference space, achieving high acceptance (70%-90%) and improved mixing.
• Normalizing flows are trained via maximum likelihood or variational inference to approximate target distributions, enhancing scalability and effective posterior estimation.

Transport Reversible Jump Markov Chain Monte Carlo (RJMCMC) methods via normalizing flows represent a significant advance for efficient Bayesian model selection and transdimensional inference. These methods utilize deep normalizing flows to construct diffeomorphic transport maps between complex model posteriors and a common reference distribution, enabling near-optimal proposals for dimension-changing jumps in RJMCMC. This approach circumvents the low acceptance rates and poor mixing inherent in classic RJMCMC by transforming transdimensional moves into near-trivial operations under the reference measure, as formalized in the works of Davies et al. (Davies et al., 2022) and subsequent extensions (Yin et al., 14 Dec 2025).

1. Mathematical Structure of Transport RJMCMC

Transport RJMCMC (TRJ) operates on a state space

$$\mathcal{X} = \bigcup_{k\in\mathcal{K}} \{ k \} \times \Theta_k,$$

where each $k$ denotes a discrete model and $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$ parameterizes model $k$. The conditional target is $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$. To facilitate dimension-matching jumps, TRJ constructs for each model $k$ a transport diffeomorphism $T_k: \Theta_k \to \mathbb{R}^{n_k}$ satisfying ideally $T_{k\sharp}\pi_k = \nu^{n_k}$ for a reference (e.g., standard normal) product measure $\nu^{n_k}$.

A transdimensional move from $(k, \theta_k)$ to $k$0 proceeds by:

1. Mapping $k$1 to the reference $k$2;
2. If $k$3, drawing $k$4; if $k$5, truncating $k$6 appropriately;
3. Establishing $k$7 through a volume-preserving pairing $k$8, e.g., concatenation or permutation in reference space;
4. Inverting $k$9 by $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$0.

This construction reduces general dimension-changing proposals to invertible, tractable mappings in reference coordinates, leveraging the expressiveness of normalizing flows to approximate the necessary transports (Davies et al., 2022, Yin et al., 14 Dec 2025).

2. Metropolis–Hastings Acceptance Criteria

The TRJ acceptance probability for a proposed move $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$1 is

$\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$2

where $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$3 is the model-jump proposal probability and $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$4 (resp. $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$5) is the reference density for auxiliary variables $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$6 (resp. $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$7).

Exact Transport Case:

If each $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$8 is an exact transport and the pairing $\theta_k\in \Theta_k \subset \mathbb{R}^{n_k}$9 preserves $k$0, the acceptance simplifies to

$k$1

With $k$2, all cross-model moves are accepted ($k$3), yielding rejection-free RJMCMC (Davies et al., 2022).

Approximate Flows:

Empirically, with trained flows, acceptance rates for cross-model moves typically exceed $k$4–$k$5 for expressive architectures such as rational-quadratic masked autoregressive flows (RQMA), far surpassing naive or affine RJMCMC (Davies et al., 2022, Yin et al., 14 Dec 2025).

3. Architecture and Training of Normalizing Flows

Flow Parameterization

• Davies et al. (Davies et al., 2022):
  • $k$8: Affine standardization (empirical mean/variance).
  • $k$9: Sigmoid/Logit transforms between $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$0 and $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$1.
  • $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$2: Stack of three masked-autoregressive spline transforms (RQMA) with rational-quadratic splines (10 bins/11 knots), conditioned by neural networks (two hidden layers, size $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$3).
• Yin & Jiao (Yin et al., 14 Dec 2025):

Uses RealNVP coupling flows for each model, with "scale" $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$4 and "translation" $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$5 functions parameterized by neural networks. Flows are either model-specific or amortized via a single conditional network with the model index as context.

Training Objective

• Maximum Likelihood:

Train $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$6 to maximize the expected likelihood of reference samples mapped by $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$7, i.e.,

$\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$8

• Variational Inference (VI):

Yin & Jiao (Yin et al., 14 Dec 2025) minimize the reverse KL divergence between the flow-based variational approximation $\pi_k(\theta_k) = \pi(k, \theta_k)/\pi(k)$9 and $k$0. Unlike pilot-run-based training, this enables fitting with samples drawn directly from the base distribution, promoting scalability and eschewing expensive MCMC pilot runs.

Amortized and Conditional Flows

For large or structured model spaces, conditional normalizing flows (CTRJ) are trained to map from a model-and-parameter-augmented reference to any model's parameterization, supporting efficient transport and dramatically reduced training cost (Davies et al., 2022, Yin et al., 14 Dec 2025).

4. Implementation Workflow and Pseudocode

A schematic pseudocode of the TRJ sampler is as follows:

$k$9 (Davies et al., 2022, Yin et al., 14 Dec 2025)

5. Empirical Performance and Benchmarks

Experimental Results

• Sinh–Arcsinh Toy Example:

With known analytic transports, TRJ achieves rejection-free moves. Approximate flows (RQMA or RealNVP) yield nearly the same performance; affine flows result in noticeably worse mixing and slower convergence of model probabilities (Davies et al., 2022, Yin et al., 14 Dec 2025).

• Bayesian Factor Analysis (6×6 covariance):

RQMA-based TRJ greatly outperforms classical Lopes–West proposals and affine flows in mixing and variance reduction (Bartolucci Bridge Estimator), offering lower RMSE in model-probability estimates at fixed computational cost (Davies et al., 2022, Yin et al., 14 Dec 2025).

• Block Variable Selection/Robust Regression:

Standard auxiliary-variable proposals nearly fail to mix between separated modes. Individual (affine/RQMA) flows and conditional flows (CTRJ) both achieve high acceptance and excellent mixing. CTRJ matches the performance of multiple separate flows at much lower training cost (Davies et al., 2022, Yin et al., 14 Dec 2025).

• TRJ-VI-NF Comparisons:

Directly minimizing the reverse KL produces maps that closely approximate the true transport, yielding higher acceptance rates, smaller variance in model probability estimates, and effective sample sizes for between-model jumps 2–3× larger than prior methods (Yin et al., 14 Dec 2025).

6. Practical Recommendations and Extensions

• Flow Design:

Choose architectures with analytic invertibility (e.g., RQMA, RealNVP). Use affine flows only when conditional posteriors are close to Gaussian.

• Preprocessing:

Standardize each training set with elementwise affine transformations computed from pilot samples to improve training efficiency and numerical stability.

• Data Requirements:

Obtain $k$1–$k$2 samples per model for supervised flow training (via MCMC/SMC pilot runs or variational flows, depending on approach).

• Model Jump Probabilities:

If marginal model probabilities $k$3 are estimable, use $k$4 to improve cross-model mixing; otherwise adopt symmetric or local proposals.

• Marginal-Likelihood Estimation:

Use the tractable flow-based variational densities for efficient estimation of $k$5 via importance sampling, enabling accurate Bayes-factor computation (Yin et al., 14 Dec 2025).

• Computational Cost:

The up-front cost of flow training is amortized by subsequent orders-of-magnitude speedups in cross-model mixing and typically dwarfed by likelihood evaluations in complex models.

• Conditional/Amortized Flows:

For large model spaces, employ amortized conditional normalizing flows (CTRJ) for scalability, training on saturated targets that pad all models to $k$6 dimensions to share statistical strength and training cost across models (Davies et al., 2022, Yin et al., 14 Dec 2025).

7. Theoretical Foundations and Extensions

“Flows for Flows” (Klein et al., 2022) introduces the nested use of normalizing flows, allowing construction of bijective maps between arbitrary distributions with known Jacobians. This principle underpins the design of deterministic RJMCMC proposals via paired flows, extending to settings with both matching and mismatched dimensions by embedding auxiliary variables and constructing invertible mappings on the joint space $k$7. This framework enables the introduction of optimal transport penalties (e.g., $k$8, Wasserstein, entropy-regularized cost) to ensure semantically meaningful transports, and supports conditional flows parameterized by model indices or side-information, which can further regularize and scale TRJ to complex, structured model spaces.

The overall methodology centralizes around replacing the design of bespoke, problem-specific RJ proposals with a learned, mathematically principled class of transport maps, ensuring detailed balance via explicit Jacobian accounting and leveraging the full expressive power of modern flow architectures.

---

References:

• Davies, Lucas, Meinshausen, "Transport Reversible Jump Proposals" (Davies et al., 2022).
• Yin, Jiao, "Transport Reversible Jump Markov Chain Monte Carlo with proposals generated by Variational Inference with Normalizing Flows" (Yin et al., 14 Dec 2025).
• Foster, Papamakarios, "Flows for Flows: Training Normalizing Flows Between Arbitrary Distributions with Maximum Likelihood Estimation" (Klein et al., 2022).