A friendly introduction to triangular transport

Measure transport methods provide a framework to transform one probability distribution into another, which is highly useful for characterizing complex target distributions by transforming them into a simpler, well-understood reference distribution.

How it works

Triangular transport is a framework to transform one probability distribution into another, allowing the characterization of a complex target distribution π by transforming it from a simpler, known reference distribution η. In generative modeling, this involves learning a transport map S that transforms reference samples z ∼ η into samples from the target distribution x = S(z) ∼ π. This operation is highly useful because it allows us to characterize conditionals π(ab∗) of the joint target distribution π(a, b) by coupling two distributions.

Key Properties of Triangular Maps

The triangular structure sets these maps apart from other measure transport methods, offering several desirable mathematical and computational properties:

1. Parsimony: The parameterization of the map function is at the user’s discretion, allowing for adjustment of complexity from resolving individual variables to overall map complexity.

2. Sparsity: Triangular maps have a natural ability to exploit conditional independence, which improves computational efficiency and robustness to smaller sample sizes and spurious correlations.

3. Numerical convenience: Constructing triangular maps involves parameterizing simple one-dimensional monotone functions, making them easy to optimize and invert.

4. Explainability: There is a clear correspondence between their constituent elements and the statistical features they represent, allowing for the description of factorizations of the target distribution using these elements.

Theoretical Foundations

The core mathematical tool is the change-of-variables formula, which relates a random variable x associated with a complicated target pdf π to a second RV z associated with a much simpler reference pdf η through an invertible, differentiable transformation S. For scalar-valued RVs, this is defined as:

π(x) = Sη(x) = η(S(x)) / ∂S(x)/∂x

This formula allows for the characterization of conditionals p (ab∗). Bayesian inference involves constructing a joint distribution p (a, b), conditioning it on a specific observation value b∗, and then normalizing the resulting slice to yield the posterior pdf p (ab∗). Triangular transport solves this challenge by first approximating an almost arbitrary joint pdf p(a, b) using measure transport and then evaluating any of its conditionals.

Sampling Conditionals and Factorization

Triangular maps naturally allow for the exploitation of conditional independence by construction. The factorization of the target distribution π in terms of marginal conditionals is given by:

π (x) = π (x1:k) S−11:k(z1:k) π (xk+1:Kx1:k) S−1k+1:K(zk+1; x1:k).

By manipulating the inversion process, one can sample conditionals of the target distribution π. If we are interested in sampling a conditional distribution p(ab∗), we define our target pdf π as the joint distribution p (a, b) and consider the manipulated samples x∗1:k to be the observations b∗ of b. The manipulated inversion samples the posterior p (ab∗).

Implementation and Optimization

The optimization problem for maps from samples seeks to minimize the Kullback–Leibler divergence between the target pdf π and its approximation Sη, i.e., Sopt = arg min S∈F D(π∥Sη). This objective function can be expanded into a Monte Carlo approximation:

J (S) = X N i1 Σ K k1 1/2 Sk(X i)2 − log ∂Sk(Xi) / ∂x k.

This objective function can be decomposed into independent objectives Jk(Sk), allowing for parallel optimization of each map component Sk. For linear separable maps, the optimization simplifies to a constrained convex problem, where coefficients bc non k are found by solving the normal equations:

bc non k = - (Pnon k⊤Pnon k + λI)−1Pnon k⊤ z> Mk,λ Pmon k cmon k.

Map Adaptation and Scalability

To address the challenge of scaling with high dimensionality, map adaptation algorithms are employed. These methods start with a simple map and iteratively propose candidate basis functions for each component Sk based on the gradient of the optimization objective function. The algorithm adds the candidate corresponding to the steepest derivative, gradually expanding complexity until a stopping criterion is met, which helps identify parsimonious maps by exploiting conditional independence. This process can be guided by graph structure learning algorithms or information criteria to help identify suitable levels of map complexity.

Toolbox and Outlook

Triangular transport methods offer parsimony, numerical convenience, and transparency for Bayesian inference.

As a fastidious researcher, I have analyzed A Friendly Introduction to Triangular Transport and identified several high-impact areas for improving Artificial Intelligence systems, particularly in fields involving uncertainty quantification, generative modeling, and Bayesian inference.

Here are the specific improvements an AI system can achieve:

1. 】Inference in High-Dimensional State-Space Models (SSMs) via Nonlinear Filtering

2. 】Enhanced Generative Modeling of Complex Target Distributions via Sample Generation

3. 】Robust Data Assimilation and Inverse Problem Solving with True Nonlinearity

4. 】Efficient Learning and Optimization of Complex, Structured Neural Networks

5. The system can perform high-fidelity Bayesian inference for complex state-space models (SSMs) by leveraging the triangular transport map to approximate intractable conditional posterior distributions, specifically in applications like those described by Marzouk et al. (2017).

6. The AI system can generate novel, high-dimensional samples from complex non-Gaussian target distributions (e.g., 400x400 pixel images of cats) by learning a transport map that transforms simple reference noise into the desired distribution, enabling robust generative modeling (Wang and Marzouk, 2022).

7. The system can solve challenging inverse problems (e.g., in geophysical or medical imaging) by using triangular transport for simulation-based inference and data assimilation, providing true nonlinear generalizations of standard filters like the ensemble Kalman filter (Spantini et al., 2022; Ramgraber et al., 2023).

8. The system can perform optimal experimental design by calculating expected information gain or mutual information using triangular transport maps, allowing for more efficient allocation of experiments in Bayesian settings (Huan et al., 2024).

9. The system can learn and deploy highly expressive, structured neural networks (maps) that are inherently sparse and robust to noise by utilizing the triangular structure property. This allows the map to naturally exploit conditional independence, significantly reducing computational complexity in high-dimensional settings (Section 3.1.3).

10. The system can perform adaptive learning of map complexity via Map Adaptation Algorithms (Baptista et al., 2023), iteratively adding basis functions with the steepest gradient to ensure the map captures necessary features without overfitting, leading to a superior bias-variance trade-off.

11. The system can generate high-quality conditional samples for complex distributions by employing Composite Maps (Section 4.1). This strategy mitigates errors from approximate maps by combining forward and inverse maps, ensuring that the resulting posterior samples accurately reflect the true target distribution, even when using simpler or less expressive transport maps.

12. The system can optimize its own map components efficiently using specialized numerical techniques like Linear Separable Maps (Appendix A). This allows for closed-form solutions for many coefficient sets, drastically speeding up the optimization process compared to general nonlinear solvers, which is crucial for real-time applications.

Decision making under uncertainty is a cross-cutting challenge in science and engineering. Most approaches to this challenge employ probabilistic representations of uncertainty. In complicated systems accessible only via data or black-box models, however, these representations are rarely known. We discuss how to characterize and manipulate such representations using triangular transport maps, which approximate any complex probability distribution as a transformation of a simple, well-understood distribution. The particular structure of triangular transport guarantees many desirable mathematical and computational properties that translate well into solving practical problems. Triangular maps are actively used for density estimation, (conditional) generative modelling, Bayesian inference, data assimilation, optimal experimental design, and related tasks. While there is ample literature on the development and theory of triangular transport methods, this manuscript provides a detailed introduction for scientists interested in employing measure transport without assuming a formal mathematical background. We build intuition for the key foundations of triangular transport, discuss many aspects of its practical implementation, and outline the frontiers of this field.