Distributed Measure Transport Maps

• The paper introduces a distributed framework for constructing invertible measure transport maps by minimizing the KL divergence under consensus constraints, enabling global feasibility through local computations.
• It employs advanced parameterizations, including polynomial chaos expansions and triangular (Knothe–Rosenblatt) maps, to manage high-dimensional and non-Gaussian distributions efficiently.
• Distributed optimization via ADMM and sequential composition strategies facilitates scalable Bayesian inference, decentralized filtering, and federated learning in multi-agent settings.

A measure transport map is an invertible transformation $T$ that pushes forward a source probability measure $P$ on a space $W\subset\mathbb{R}^D$ to a target measure $Q$, i.e., $T_\# P = Q$. Distributed construction of such maps seeks to realize $T$ through parallel or decentralized computation, leveraging either data parallelism or networked agents while maintaining consistency and feasibility. This is motivated by large-scale Bayesian inference, decentralized filtering, federated learning, and multi-agent settings where full data aggregation is impractical or undesirable. The core methodological themes integrate convex optimization, polynomial chaos or triangular (Knothe–Rosenblatt) parameterizations, consensus algorithms such as the alternating direction method of multipliers (ADMM), and sequential composition strategies for high-dimensional or non-Gaussian distributions (Mesa et al., 2018, Last et al., 2021, Grange et al., 2023).

1. Mathematical Formulation and Problem Setting

The distributed construction of measure transport maps formalizes the problem as finding an invertible, orientation-preserving diffeomorphism $T:W\to W$ such that $T_\# P=Q$. The change-of-variable formula reads $p(u) = q(T(u)) \det JT(u)$, where $JT$ is the Jacobian of $P$0. A prevalent variational approach minimizes the Kullback–Leibler (KL) divergence:

$P$1

subject to $P$2 almost everywhere. Under the log-concavity of $P$3 and mild regularity of $P$4, the underlying variational objective is convex with unique minimizer (Mesa et al., 2018). In random measure settings, one generalizes to stationary and ergodic random measures $P$5 on $P$6, with allocations (transport maps) defined to satisfy translation invariance and mass-balance (Last et al., 2021).

2. Parameterizations and Locality

Two primary structural parameterizations prevail:

• Polynomial Chaos Expansion: The map $P$7 is componentwise represented via a truncated multivariate polynomial chaos basis:

$P$8

where $P$9 collects $W\subset\mathbb{R}^D$0 basis functions and $W\subset\mathbb{R}^D$1. The parameter $W\subset\mathbb{R}^D$2 grows polynomially with $W\subset\mathbb{R}^D$3 and the maximal basis order, but is mitigated via single-univariate or Knothe–Rosenblatt (KR) structures (Mesa et al., 2018).

• Triangular (KR)-type Maps: The map has lower-triangular structure

$W\subset\mathbb{R}^D$4

with monotonicity enforced via positivity constraints on diagonal parameters. This structure aligns with conditional sampling and dimension reduction, as each component depends only on its predecessors (Grange et al., 2023).

These parameterizations enable local updates, as each component or block can be computed and updated with local data or samples, subject to consensus constraints.

3. Distributed Optimization and Consensus Algorithms

Distributed computation is achieved via consensus-constrained optimization frameworks. Given sample-based empirical risk objectives, parameter blocks $W\subset\mathbb{R}^D$5 or $W\subset\mathbb{R}^D$6 associated with each worker or node are constrained to agree globally, e.g.,

$W\subset\mathbb{R}^D$7

ADMM solves this via alternating updates:

• Local variable update by minimizing an augmented Lagrangian (often convex subproblems per data block).
• Global consensus update, typically averaging local parameters plus dual variables.
• Dual variable update to enforce equality constraints.

Communication is limited to exchanging parameter vectors and locally-computed statistics. Under convexity and regularity, ADMM converges at $W\subset\mathbb{R}^D$8 rates in objective and feasibility (Mesa et al., 2018, Grange et al., 2023).

In the context of sensor networks or multi-agent settings, consensus-ADMM enables each node to operate autonomously with its neighborhood information, with periodic synchronization across the network (Grange et al., 2023).

4. Sequential Composition for High-Dimensionality

To address the complexity of learning a highly nonlinear $W\subset\mathbb{R}^D$9 in large or non-Gaussian systems, sequential composition decomposes $Q$0 as a series of simpler maps:

$Q$1

where each $Q$2 evolves the source measure incrementally toward the target. This draws on the Jordan–Kinderlehrer–Otto (JKO) interpretation of gradient flows in the space of probability measures:

$Q$3

with each step approximated by solving a map-based optimization. Each component map is close to identity and can be parameterized with lower-order polynomial chaos or triangular representations. Only 10–20 such maps are usually required, even for $Q$4, as each corrects only residual discrepancies (Mesa et al., 2018).

5. Distributed Measure Transport for Random Measures

For jointly stationary and ergodic random measures $Q$5 on $Q$6 of equal intensity, translation-invariant allocations may be constructed using distributed algorithms. When the target has atoms, a stable marriage—type algorithm based on spatial Gale–Shapley dynamics grows atomic "claim-regions" in balls, incrementally assigning mass until exact matching is achieved (Last et al., 2021). If both source and target are diffuse, an auxiliary point process $Q$7 is required for construction:

• Pair $Q$8 and $Q$9 to $T_\# P = Q$0 via discrete allocations.
• Within each pair, perform quantile-based matching via measurable bijections and local inversion.

Mass-balance is guaranteed by translation equivariance and ergodic arguments, with allocations satisfying

$T_\# P = Q$1

The construction is fully local: each site requires only finite-range knowledge, making the process scalable and inherently distributed (Last et al., 2021).

6. Computational Complexity, Scalability, and Limitations

Resource requirements are determined by the dimension $T_\# P = Q$2, polynomial/statistical basis size $T_\# P = Q$3, and the number of parallel workers or nodes $T_\# P = Q$4:

• Computation: Each worker’s per-iteration complexity is $T_\# P = Q$5 for evaluating basis functions and their derivatives; additional cost is determined by local optimization or eigen-decompositions (the latter reduced in KR-case).
• Communication: Each iteration communicates $T_\# P = Q$6 (parameter matrices) and $T_\# P = Q$7 (for Jacobian- or moment-based structures) per worker.
• Scaling: Memory constraints for basis expansion can become prohibitive for large $T_\# P = Q$8, though single-univariate or triangular forms mitigate this growth to $T_\# P = Q$9 per map.
• Empirical results: For $T$0 samples and $T$1, 8-GPU clusters allow for near-linear weak scaling (Mesa et al., 2018), with feasible multi-block distributed training up to $T$2.
• Limitations: Target measure $T$3 must be log-concave for convexity guarantees. Polynomial chaos bases, even in KR-variant, do not eliminate but reduce dimensional scaling. Inversion (for backward sampling) can become a bottleneck for $T$4. These constraints motivate research into advanced bases and asynchronous or stochastic optimization (Mesa et al., 2018).

7. Applications and Broader Impacts

Applications for distributed construction of measure transport maps include:

• Bayesian inference: Posterior generation for high-dimensional models, demonstrated in LASSO-type regression and Gaussian process models. Posterior samples are generated deterministically and without MCMC autocorrelation (Mesa et al., 2018).
• Generative modeling: Nonlinear generative flows for image data (MNIST), where one-to-one normalizing flow maps generate samples via distributed training procedures.
• Distributed filtering: Nonlinear filtering in sensor networks using triangular maps fused via consensus algorithms; this avoids global data aggregation and allows for privacy-respecting distributed estimation (Grange et al., 2023).
• Multi-agent and federated optimization: Distributed barycenter computation, decentralized generative modeling, and federated Bayesian inference, where privacy constraints prioritize local computation and parameter sharing over raw data exchange.

The synthesized framework enables scalable, parallelized, and privacy-aware computation for statistical and probabilistic modeling across data-intensive domains (Mesa et al., 2018, Grange et al., 2023, Last et al., 2021).