On the representation and learning of monotone triangular transport maps

Many sampling, estimation, and inference algorithms seek to characterize a somehow intractable or complex probability distribution $\bm{\mu}$ on $\mathbb{R}^{d}$. Transportation of measure provides a useful and versatile approach to this problem. The underlying idea is to construct a coupling of $\bm{\mu}$ with a tractable “reference” distribution $\bm{\nu}$ on $\mathbb{R}^{d}$—for instance, a standard normal. Formally, one jointly constructs a pair of random variables $(\bm{X},\bm{Z})$ such that $\bm{X}\sim\bm{\mu}$ and $\bm{Z}\sim\bm{\nu}$. A special class of couplings is given by deterministic transformations $S\colon\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ such that $S(\bm{X})=\bm{Z}$ in distribution, and a transformation $S$ that satisfies this property is called a transport map [69]. As a result, if the transport map is invertible, one can generate realizations of $\bm{X}$ by first simulating $\bm{Z}$. If $\bm{\mu}$ and $\bm{\nu}$ have densities $\pi$ and $\eta$ with respect to a common base measure, one can also explicitly represent the target density $\pi$ as a transformation of the reference density $\eta$. The construction of such transport maps has found numerous applications: density estimation [62, 2, 18, 21]; variational Bayesian inference [20, 50, 10]; generative modeling of images, video, and other structured objects [42, 27]; likelihood-free inference [43, 47, 33]; and beyond.

In general, there exist infinitely many transport maps between two absolutely continuous distributions on $\mathbb{R}^{d}$. In this paper, we focus on a specific, canonical choice: triangular transport maps [8] of the form

where each component function $S_{k}$ depends only on the $k$ variables $\bm{x}_{\leq k}\coloneqq(x_{1},\ldots,x_{k})$ and is monotone increasing with respect to the last input $x_{k}$. In particular, each function $S_{k}$ encodes an increasing rearrangement [69] between ordered marginal conditionals of $\bm{\mu}$ and $\bm{\nu}$, i.e., $\bm{\mu}(\mathrm{d}x_{k}|\bm{x}_{<k})$ and $\bm{\nu}(\mathrm{d}z_{k}|\bm{z}_{<k})$. Later we will discuss properties of such triangular maps—known also as Knothe–Rosenblatt (KR) rearrangements [69, 51, 53]—more precisely, but we comment here on their utility. First of all, triangular structure facilitates computational tractability: $S$ is easy to invert and the determinant of its Jacobian (a lower triangular matrix) is easy to evaluate (see Section 2). Triangular maps have thus been used extensively in the density estimation, inference, and generative modeling applications noted above. For example, triangular maps are the building blocks of many normalizing flows, popularized by the machine learning community [28, 44]; specifically, autoregressive normalizing flows define triangular maps [25] via particular structural choices and parameterizations, and compose these maps to produce more expressive transformations. More fundamentally, because triangular maps expose certain conditionals of $\bm{\mu}$, they are particularly well suited to conditional density estimation and conditional sampling; we will describe this link in Section 2. Triangular maps also inherit sparsity from the conditional independence properties of $\bm{\mu}$ and $\bm{\nu}$, as detailed in [61].

It is worth noting, of course, that other canonical choices of transport have different attractive features. For instance, optimal transport maps are invariant under relabeling of the coordinates or more general isometries on $\mathbb{R}^{d}$ (unlike triangular maps), and have deep links to partial differential equations [69]. But optimal transport maps are in general more challenging to represent, evaluate, and estimate from data in the continuous setting; moreover, they do not enjoy such a direct link to conditioning or to graphical models.

Many representations and finite-dimensional parameterizations of monotone triangular maps have been proposed in recent years. These include representations based on polynomials [35, 25], radial basis functions [62], neural networks of varying capacity [18, 27], and tensor decompositions [15, 16]. A core challenge in this setting is to satisfy the monotonicity constraint $\partial_{x_{k}}S_{k}>0$. For instance, one might enforce the monotonicity constraint at a finite collection of points in the support of $\pi$ [46], but this approach cannot in general guarantee that $S$ is monotone over the entire support of $\pi$. Other efforts have sought to enforce monotonicity by construction—via the parameterization of $S$ itself. For example, [45] employs map components with affine dependence on the last variable, i.e., $S_{k}(\bm{x}_{\leq k})=\alpha(\bm{x}_{<k})+\exp(\beta(\bm{x}_{<k}))x_{k}$, where $\alpha$ and $\beta$ are neural networks. While $S$ is then guaranteed to be monotone, it can only represent a restricted class of distributions $\bm{\mu}$. (If $\bm{\nu}$ is Gaussian, then $\bm{\mu}$ can only be a product of Gaussian marginal conditionals.) Such representations therefore cannot consistently approximate the true KR rearrangement for general $\bm{\mu}$. Recent work [63] has shown that a composition of such affine maps, interleaved with rotations and permutations, can approximate a general class of distributions, though approximation rates remain unknown. The required rotations or permutations break the overall triangular structure of the transformation, however. Alternatively, to increase the “expressiveness” of a given triangular function, [72, 22, 19] have introduced more general parametric representations of the monotone component functions $S_{k}$. For distributions with analytic densities on bounded domains, a complete approximation theory for the KR map was recently developed in [75, 76], using polynomial or ReLU neural network representations of range-constrained monotone triangular functions.

Despite these myriad proposals, relatively little attention has been paid to the structure and tractability of the optimization problem involved in learning triangular maps (e.g., in estimating a map given an i.i.d. sample $\{\bm{X}^{i}\}$ from $\bm{\mu}$). Properties of this optimization problem are intimately tied to the representation and parameterization of the associated map. It is desirable to have a flexible and general representation—one that can consistently recover the KR rearrangement for a broad class of distributions $(\bm{\mu},\bm{\nu})$—that at the same time makes optimization tractable. It is also desirable to have adaptivity: a parameterization whose size or complexity can be adapted to properties of the target distribution and the available sample size, for good empirical statistical performance.

This paper directly addresses these desiderata. We do so by developing and analyzing a functional framework for representing and learning triangular maps. Our main contributions are as follows. We propose a general representation of monotone triangular functions, based on a rectification operator $\mathcal{R}_{k}$ that transforms sufficiently smooth non-monotone functions $f_{k}:\mathbb{R}^{k}\to\mathbb{R}$ into monotone component functions $S_{k}$ of a triangular map. This operator takes the form

where $g\colon\mathbb{R}\to\mathbb{R}_{>0}$ is a positive function. We then analyze the infinite-dimensional optimization problem associated with learning maps from data, recast as optimization over functions $\{f_{k}\}_{k=1}^{d}$. Specifically, we establish conditions on the rectification operator and on the target distribution such that the resulting optimization problem is well-posed, smooth, and has no spurious local minima. Under further conditions on the target distribution (essentially that it has Gaussian tails), we show that the optimization problem has a unique global minimizer corresponding to the KR map.

These theoretical results guarantee, in practice, fast and reliable learning of monotone triangular maps given an appropriate function space $V_{k}$ for each $f_{k}$. The second main contribution of our paper is algorithmic: given a hierarchical basis (e.g., polynomials or wavelets) for each $V_{k}$, we propose a greedy adaptive procedure to learn parametric representations of $f_{k}$. The procedure naturally produces map representations that are sparse and interpretable—in particular, it exploits and implicitly discovers conditional independence. We use these learned maps for density estimation, given an i.i.d. sample $\{\bm{X}^{i}\}$ from $\pi$. Maintaining a strict triangular structure also exposes marginal conditionals of the target density, thus enabling conditional density estimation. Our numerical experiments show that the algorithm provides robust performance at small-to-moderate sample sizes, and constitutes a semi-parametric approach that naturally links map complexity to the size of the data.

The remainder of the paper is organized as follows. Section 2 recalls properties of triangular transport maps and introduces some estimation problems of interest. Our main theoretical contributions are in Section 3, which introduces a framework for representing monotone triangular maps and analyzes the resulting optimization problem. Section 4 then introduces our greedy adaptive procedure for learning maps, and Section 5 contains numerical experiments. Proofs of certain theoretical results are deferred to the appendix.

Consider the unsupervised learning problem of approximating a target probability density function $\pi$ defined on $\mathbb{R}^{d}$, given an i.i.d. sample from $\pi$. Our goal is to construct a sufficiently smooth and invertible map $S\colon\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ such that the pullback density

is a good approximation to $\pi$, where $\eta$ is a simple/tractable probability density function on $\mathbb{R}^{d}$. The choice of $\eta$ is a degree of freedom of the method, and here we take $\eta(\bm{x})\propto\exp(-\|\bm{x}\|^{2}/2)$; i.e., $\eta$ is the density of the standard normal distribution on $\mathbb{R}^{d}$, where $\|\cdot\|$ is the canonical norm of $\mathbb{R}^{d}$.

To ensure invertibility of $S$, a common practice is to constrain $S$ to be an increasing lower triangular map of the form (1), where each component $S_{k}\colon\mathbb{R}^{k}\rightarrow\mathbb{R}$ is such that $x_{k}\mapsto S_{k}(\bm{x}_{<k},x_{k})$ is increasing for all $\bm{x}_{<k}=(x_{1},\ldots,x_{k-1})\in\mathbb{R}^{k-1}$. Such a map is easy to invert¹¹1For any $\bm{z}\in\mathbb{R}^{d}$, $\bm{x}=S^{-1}(\bm{z})$ can be computed recursively as $x_{k}=T^{k}(\bm{x}_{<k},z_{k})$ for $k=1,\dots,d$, where the function $T^{k}(\bm{x}_{<k},\cdot)$ is the inverse of $x_{k}\mapsto S_{k}(\bm{x}_{<k},x_{k})$. In practice, evaluating $T^{k}$ requires solving a root-finding problem which is guaranteed to have a unique (real) root, and for which the bisection method converges geometrically fast. Therefore, $S^{-1}(\bm{z})$ can be evaluated to machine precision in negligible computational time. and has a lower triangular Jacobian $\nabla S(\bm{x})$ so that $|\det\nabla S(\bm{x})|=\prod_{k=1}^{d}\partial_{k}S_{k}(\bm{x}_{\leq k})$ is readily computable. This is a major benefit of working with monotone triangular maps. This structure in fact corresponds to the Knothe–Rosenblatt (KR) rearrangement $S_{\text{KR}}$: the increasing lower triangular map satisfying

For a measure $\bm{\mu}$ on $\mathbb{R}^{d}$ that is absolutely continuous with respect to a Gaussian measure $\bm{\nu}$ (and hence has a density $\pi$ with respect to the Lebesgue measure), the KR rearrangement $S_{\textrm{KR}}$ exists and is the unique map of the form (1) that pulls back $\eta$ to $\pi$ (or equivalently pushes forward $\pi$ to $\eta$), up to sets of measure zero [8].

A useful way to measure the difference between $\pi$ and its approximation $S^{\sharp}\eta$ is the Kullback–Leibler (KL) divergence $\mathcal{D}_{\text{KL}}(\pi||S^{\sharp}\eta)=\int\log(\pi/S^{\sharp}\eta)\mathrm{d}\pi$. As explained below, this choice has direct links to maximum likelihood estimation of $S$. Furthermore, the following inequality shows that convergence in the KL sense $\mathcal{D}_{\text{KL}}(\pi||S^{\sharp}\eta)\rightarrow 0$ implies convergence of $S$ towards $S_{\mathrm{KR}}$ in the $L^{2}_{\pi}$ sense. This result is a direct consequence of Corollary 3.10 in [8]; see Section A.1 for a proof.

Since the standard normal density $\eta$ is a product of its marginal densities, the KL divergence decomposes as

where the functionals $\mathcal{J}_{1},\ldots,\mathcal{J}_{d}$ are given by

Minimizing the KL divergence (5) over triangular maps $S$ of the form (1) is therefore equivalent to independently minimizing each objective functional $\mathcal{J}_{k}$ to find the associated map component $S_{k}$ [35]. Solution of these optimization problems is thus embarrassingly parallel. This parallel structure was also exploited for Cholesky factorization via KL minimization in [54]. In addition, minimizing each objective $s\mapsto\mathcal{J}_{k}(s)$ over functions $s\colon\mathbb{R}^{k}\rightarrow\mathbb{R}$ that are strictly increasing in the last variable is a strictly convex optimization problem; see Lemma 10.

Given an i.i.d. sample $\{\mathbf{X}^{i}\}_{i=1}^{n}$ from $\pi$, we can replace the expectation in (6) by the sample average, which yields the objective

Minimizing (7) under the constraint $\partial_{k}s(\bm{x}_{\leq k})>0$ produces an estimator $\widehat{S}_{k}$ of $S_{\text{KR},k}$. The collection of all such component functions defines an estimator $\widehat{S}=(\widehat{S}_{1},\ldots,\widehat{S}_{d})$ of the KR rearrangement, along with an estimate of the density $\pi$ as $\widehat{\pi}(\bm{x})\coloneqq\widehat{S}^{\sharp}\eta(\bm{x})$. This $\widehat{S}$ is also the maximum likelihood estimator of $S_{\text{KR}}$, i.e.,

with the optimization being over the space of monotone increasing and triangular maps of the form (1). This connection between maximum likelihood estimation and minimization of an empirical forward KL divergence is standard. Furthermore, convexity of the optimization objective is preserved when replacing the expectation in (6) by the sample average in (7).

A core question when estimating maps and densities by minimizing (7) is how to parameterize sufficiently expressive monotone map components $S_{k}$—i.e., maps capable of representing a wide class of distributions $\pi$—while ensuring that the optimization problem can be solved efficiently. As explained in the introduction, this question is intimately tied to the monotonicity constraint $\partial_{k}S_{k}(\bm{x}_{\leq k})>0$. For example, choosing a linear parameterization for $S_{k}$ that admits only affine dependence on $x_{k}$ (to easily enforce monotonicity) allows the map component be identified efficiently through the solution of a least-squares problem (see [60, Appendix A]), but such maps can only capture distributions that factor into a sequence of Gaussian marginal conditionals. On the other hand, a more complex ansatz for $S_{k}$ will often yield a much more difficult optimization problem. Note that with any nonlinear parameterization of $S_{k}$, the convexity of (6) and (7) with respect to $S_{k}$ does not in general yield convexity in the parameters. As we shall demonstrate later, many nonlinear parameterizations that enforce monotonicity yield optimization problems that may not even be smooth, and that have many local minima. We will address these issues in Section 3.

In this section, we define a general representation for components of monotone triangular maps and present our main theoretical results. The essential idea, as mentioned in the introduction, is to express each component $S_{k}$ as a nonlinear transformation $S_{k}=\mathcal{R}_{k}(f)$ of a smooth function $f$, where $\mathcal{R}_{k}$ (2) is an operator that enforces the monotonicity constraint by construction. We then identify the map component by solving the re-parameterized optimization problem

where $V_{k}$ is a linear space of functions in which we seek $f$. With this nonlinear transformation, we lose the convexity of the constrained problem $\min_{\{s\colon\partial_{k}s>0\}}\mathcal{J}_{k}(s)$, but obtain an unconstrained minimization problem instead. It turns out that, with appropriate conditions on $\mathcal{R}_{k}$, the transformed optimization problem retains many desirable properties.

In Section 3.1, we discuss the construction of $\mathcal{R}_{k}$ and motivate certain critical choices therein. Then, in Section 3.2, we analyze the regularity of the composed objective functional $\mathcal{L}_{k}$ for a certain choice of the function space $V_{k}$. In Section 3.3 we discuss the existence and uniqueness of solutions to (10), and describe conditions under which solving (10) exactly recovers the KR rearrangement.