Constructive conditional normalizing flows

TL;DR

This work constructs explicit, flow-based transports that approximate a target diffeomorphism $\upphi$ and the resulting pushforward $\upphi_{\#}\upmu$ within controlled errors. It hinges on a constructive pipeline: approximate $\upphi$ by a Lagrange interpolant, factor it into measure-preserving and compressible components via a polar-like decomposition, and realize each factor as neural-ODE flows or permutation-based swap flows, with stability ensuring TV closeness of the pushforwards. Two complementary results are proved: a general theorem yielding $L^p$-approximation and TV control with finitely many switches (scaling at best like $\varepsilon^{-d}$ in general), and a second, Maurey-inspired construction giving $1/\varepsilon^{2}$ switch costs for smoother maps such as Knöthe–Rosenblatt or Barron-represented flows. The paper also validates the Knöthe–Rosenblatt rearrangement within this framework and discusses the role of activation choices (e.g., ReLU) in preserving exact log-Jacobian increments, highlighting practical implications for conditional sampling and measure-preserving transport.

Abstract

Motivated by applications in conditional sampling, given a probability measure $μ$ and a diffeomorphism $φ$, we consider the problem of simultaneously approximating $φ$ and the pushforward $φ_{\#}μ$ by means of the flow of a continuity equation whose velocity field is a perceptron neural network with piecewise constant weights. We provide an explicit construction based on a polar-like decomposition of the Lagrange interpolant of $φ$. The latter involves a compressible component, given by the gradient of a particular convex function, which can be realized exactly, and an incompressible component, which -- after approximating via permutations -- can be implemented through shear flows intrinsic to the continuity equation. For more regular maps $φ$ -- such as the Knöthe-Rosenblatt rearrangement -- we provide an alternative, probabilistic construction inspired by the Maurey empirical method, in which the number of discontinuities in the weights doesn't scale inversely with the ambient dimension.

Figures (7)

• Figure 1: Let $G=S_3$, where $(1,2,3)$ is the identity. Take the symmetric generating set $S=\{(2,1,3),(1,3,2)\}$ (the transpositions $(12)$ and $(23)$ in cycle notation), so $S^{-1}=S$. The Cayley graph $\mathrm{Cay}(G,S)$ has vertex set $G$ and an (undirected) edge between $g_1,g_2$ iff $g_2=sg_1$ for some $s\in S$. We color the edge blue if $g_2=(1,3,2)g_1$ and green if $g_2=(2,1,3)\,g_1$. The resulting graph is a $2$-regular connected graph on $6$ vertices, hence a single $6$-cycle $C_6$, and its diameter is $3$.
• Figure 2: \ref{['thm: gen.thm']} allows us to approximate, in particular, the optimal transport map $\upphi$ between $\upmu$ and a target absolutely continuous measure $\upnu$ by an approximate map $\upphi_\varepsilon$. However, even if the densities $\upphi_{\#}\upmu$ and $\upphi_{\varepsilon\#}\upmu$ are close in $\mathsf{TV}$ and the maps are close, the trajectory $t\mapsto\upmu^\varepsilon(t)$ given by \ref{['eq: cont.eq']} is not close to $((1-t)\mathsf{id}+t\upphi)_\#\upmu$ for all times $t$.
• Figure 3: A triangulation and its image by a piecewise affine map $\upphi_{\mathscr{L}}$.
• Figure 4: The decomposition of $\upphi_{\mathscr{L}}$ per \ref{['prop: lagrange.decomposition']}.
• Figure 5: \ref{['lem: mp.neural.ode.swap']}: the measure preserving map $m$ swaps the colored squares and leaves the whites invariant.

Theorems & Definitions (38)

• Theorem 1.1
• Theorem 1.2
• Definition 1.3
• Proposition 1.4
• Remark 1.5: Width
• Remark 1.6: Geodesics
• Theorem 2.1: iwaniec2017triangulation
• Remark 2.2
• Proposition 2.3
• proof