Learning the Structure of Triangular Transport Maps

Morten Blørstad; Pekka Parviainen; Berent Ånund Strømnes Lunde

cs.LGSep 29, 2026

Learning the Structure of Triangular Transport Maps

Authors: Morten Blørstad, Pekka Parviainen, Berent Ånund Strømnes Lunde

Organizations: Department of Informatics, University of Bergen · Equinor

Abstract

Triangular transport maps provide a flexible approach to sampling-based probabilistic modeling, including density estimation, generative modeling, and Bayesian inference. They transform an unknown target distribution into a simpler reference through a monotone triangular map. The map structure is defined by a variable ordering and sparsity pattern, which together encode a directed acyclic graph. Map quality can depend strongly on this structure, yet finding a good structure is computationally expensive because each candidate generally requires fitting a different map. A central challenge is therefore to learn density and structure jointly, while keeping computation manageable as dimension grows. We introduce Self-Structuring Transport Maps (SSTM), which learn the map, ordering, and sparsity jointly. We use SoftSort to learn the variable ordering and $L_0$ gates to learn the sparsity, while preserving a triangular structure. To keep the map scalable, we use a monotone BatchEnsemble that shares one weight matrix across all map components through rank-one adapters. Across synthetic and real data, jointly learning the structure and map gives better density estimates than estimating the structure first. When the structure is identifiable from the density, SSTM matches the density performance of a map fitted with the true structure and outperforms autoregressive flows. On large datasets, SSTM is competitive with autoregressive flows.

Figures & tables

Panel of Figure 1 — Figure 1: Overview of SSTM . The forward map pushes $\mathbf{x}\sim\pi$ to $\mathbf{z}\sim\eta$ , while the inverse uses sequential bisection. The learned ordering and sparsity gates form the sparse triangular mask $\tilde{M}=g\odot M$ , which is applied to the input. Each of the $K$ masked rows is input to one map component $T_{k}$ . The $K$ components share weights $\mathbf{W}$ and use rank-one adapters $(\mathbf{r}_{k},\mathbf{s}_{k},\mathbf{b}_{k})$ . The objective $\mathcal{L}=\mathcal{J}(T)+\mathcal{L}_{0}$ jointly optimizes the map, ordering, and sparsity. Cell colours: orange, diagonal variable; blue, conditioning variable; white, masked off.

Figure 2: The four graph families. One generated DAG example per family at

K=10

(edges oriented parent-to-child;

|E|

= edge count). A : Erdős–Rényi, random edges. B : scale-free, a few high-degree hubs. C : Neal’s funnel, a star with one root parenting all others. D : two-level hierarchical, roots feed a middle layer that generates the leaves.

Figure 3: Gap to the true density per data-generating process (columns) and dimension

K

(rows). Per-curve bands are

\pm 1

standard deviation over five runs.

	Power	Gas	Hepmass	Miniboone
	$K=6$	$K=8$	$K=21$	$K=43$
SSTM (ours)	$-0.527\pm 0.020$	$-10.854\pm 0.039$	$16.153\pm 0.312$	$9.925\pm 0.137$
MAF	$-0.454\pm 0.011$	$-11.957\pm 0.033$	$17.058\pm 0.250$	$10.314\pm 0.100$
NSF	$-0.650\pm 0.003$	$-13.031\pm 0.023$	$14.268\pm 0.262$	$10.300\pm 0.089$
GNF ⋆	$-0.62\pm 0.04$	$-10.15\pm 0.15$	$14.17\pm 0.13$	$16.23\pm 0.52$

Table 1: Real-data density estimation on the UCI benchmarks, Mean NLL

\pm

one standard deviation over five runs. Results followed by a star are copied from the Wehenkel and Louppe (2021) which reports over 3 runs.

Figure 4: Density estimation on the Sachs data: test NLL on the pooled ( A ) and observational ( B ) subsets. Small markers are single runs, bars the mean

\pm 1

standard deviation over five runs. Values beyond the axis range are marked at the frame with their value.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Data
	Power	Gas	Hepmass	Miniboone
Dimension $K$	6	8	21	43
Training points	1,659,917	852,174	315,123	29,556
Validation points	184,435	94,685	35,013	3,284
Test points	204,928	105,206	174,987	3,648
Common (all models)

Appendix

Table 2: Hyperparameters of the density experiment on the UCI benchmarks, per dataset. Data: the standard splits of Papamakarios et al. [23] . Common: the per-dataset batch size and step budget of Durkan et al. [10, Table 5] , converted to whole epochs and shared by all models. Flows: their published per-dataset settings. SSTM : per-dataset size, chosen on validation.

Figure 5: Structure recovery on the synthetic data at

N=1000

, per data-generating process (columns) and dimension

K

(rows). Metrics are

[0,1]

with

1

best; SHD is inverted and normalized. Lines are means, bands the min to max over five runs.

Figure 6: Structure recovery on the synthetic data at

N=200

, per data-generating process (columns) and dimension

K

(rows). Metrics are

[0,1]

with

1

best; SHD is inverted and normalized. Lines are means, bands the min to max over five runs.

Figure 7: Structure recovery on the Sachs data against the 20-edge reference graph, on the pooled ( A ) and observational ( B ) subsets. Metrics are scaled to

[0,1]

with

1

best; SHD is inverted and normalized. Lines are means, bands the min to max over five runs.

Figure 8: Method ablation: each variant changes one design choice of the default configuration and is compared against the default per dataset, dimension

K

, and training-set size

N

. Each marker is one such setting, averaged over five seeds (color = dataset, shape =

K

, fill =

N

). The black line is the median over settings and the grey band the interquartile range. The dashed line marks the default. A : gap to the true density relative to the default (ratio, log scale,

1

= default, right = worse). B : directed

F_{1}

relative to the default (difference,

0

= default, right = better).

Flow Map Learning

Flow map learning trains neural models to predict finite-time state or distribution transitions directly, enabling fast simulation, control, and generation.

See all 49 papers

The Fundamental Limits of Valid Transport Map Estimation
Jun 29, 2026Sivaraman BalakrishnanMinimax Estimation Generative Modeling
Department of Statistics and Data Science† · Machine Learning Department‡
Beckmann Transport Models: From Autonomous Flows to One-Step Maps
Aug 3, 2026Lee Cheuk-Kit, Florentin Coeurdoux, Yuyuan Chen +5Flow Matching One-Step Generative Modeling
1Harvard University · 2Capital Fund Management · University of Pennsylvania +2
On the Error Analysis of Triangular Transport Maps for Conditional Generation via MMD-based Objectives
Oct 24, 2026Michele Martino, Bamdad Hosseini, Amir Sagiv +1Conditional Generative Modeling Conditional Distribution Estimation
Department of Applied Mathematics, University of Washington, Seattle, WA, USA · Department of Mathematical Sciences, New Jersey Institute of Technology, Newark, NJ, USA · Department of Aeronautics and Astronautics, University of Washington, Seattle, WA, USA

Jun 29, 2026J/K move · Enter open · S save

The Fundamental Limits of Valid Transport Map Estimation

Sivaraman Balakrishnan

Department of Statistics and Data Science† · Machine Learning Department‡

Minimax Estimation Generative Modeling Optimal Transport

Open paper AskPDF arXiv

Many modern generative modeling methods, including diffusion models, normalizing flows, and flow matching, estimate transport maps or plans between distributions without explicitly targeting an optimal transport (OT) map. In applications like generative modeling, the transport cost itself is irrelevant, and this makes it natural to target maps which are more tractable from either a statistical or computational standpoint. In this short note, we formalize the task of estimating any valid transport map in a rigorous minimax framework. One consequence of this framing is that it yields sample complexity lower bounds for any method whose learned object is evaluated as a transport map or plan, including flow matching and diffusion-based generative models, in settings where direct analysis would be challenging due to the analytic complexity of the methods and their target maps. We observe that, under standard, though strong, stability assumptions from the OT literature, estimating any valid transport map is statistically as hard as estimating the OT map. We complement these results with some examples showing that when these stability assumptions fail, alternative transport maps can be learned substantially more accurately than the OT map. Our minimax framing provides a rigorous foundation for understanding the statistical limits of modern transport-based generative methods and clarifies when targeting sub-optimal maps can provide real statistical advantages.

Abstract

Flow Map Learning

Explore similar work

The Fundamental Limits of Valid Transport Map Estimation

Beckmann Transport Models: From Autonomous Flows to One-Step Maps

On the Error Analysis of Triangular Transport Maps for Conditional Generation via MMD-based Objectives