REVIEW 2 major objections 8 minor 1 cited by
Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows
T0 review · 2 major / 8 minor · reviewed 2026-08-07 · deepseek-v4-flash
Pith's one-line read The paper claims that Wasserstein-2 barycenters of absolutely continuous distributions can be computed as the pushforward of a weighted average of conditional normalizing flow maps trained by likelihood maximization plus a decreasing…
desk verdict A practical and honest algorithm for Wasserstein barycenters with conditional flows, whose stated theoretical guarantee is compromised by the unproven decreasing-weight schedule. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is a conditional normalizing flow, an invertible neural map f(z,s) trained so that for each condition s it pushes a common Gaussian latent onto the input distribution μ_s. The load-bearing identity is Theorem 3.3: for any minimizer f of the expected conditional variance objective, the weighted average h(z)=Σ_s w_s f(z,s) pushes the latent onto the Wasserstein-2 barycenter. Invertibility does double duty: it supplies the bijective maps needed for the L^p reformulation of optimal transport, and it provides the inverse passes that return optimal transport maps between inputs and the barycenter without further optimization.
What would settle it
On a Gaussian location-scatter family where the optimal transport map is known in closed form, train Algorithm 1 for increasing numbers of iterations and measure the L2-UVP of the recovered map against that closed form; the paper's claim implies the error should converge to near zero, so a persistent plateau above a small tolerance would show the training objective is not reaching the optimum.
Extended reading notes
Core claim
Under mild assumptions (all measures absolutely continuous with finite second moment, weights summing to one), the paper proves two structural identities. Theorem 3.2 says that for any absolutely continuous latent, the p-Wasserstein cost between two distributions equals the minimum over bijective pushforward pairs of the L^p distance between their maps, and that composing one map with the inverse of the other yields the optimal transport map. Theorem 3.3 says the weighted Wasserstein-2 barycenter is the law of a weighted sum of the slice maps, where the slice maps minimize an expected conditional variance among all bijective pushforwards onto the input distributions. The algorithms implement these theorems by training a conditional normalizing flow with log-likelihood plus a decreasing L2 cost; inverting the flow and passing through the weighted sum gives the optimal transport map from each input to the barycenter, and the pushforward of the latent under the weighted sum is a generative model of the barycenter.
Load-bearing premise
The whole method depends on the trained flow being flexible enough to represent the exact optimal transport maps from the shared latent to each input distribution, and on the decreasing weight on the transport cost steering training to a point where those map constraints and the cost are both satisfied; the authors note that well-conditioned flows are not known to approximate every bijection.
Editorial extensions
If this is right
- Optimal transport maps from every input distribution to the barycenter are obtained for free as x↦h(f_θ^{-1}(x,s)), with no additional computation.
- The barycenter has a direct generative model h#λ, so new samples can be drawn without querying the input distributions.
- No minimax or bi-level adversarial loop is required; both the likelihood and the transport cost are optimized with ordinary gradient descent.
- The conditional architecture amortizes across conditions, so one flow serves all inputs and training scales to hundreds of input distributions, as demonstrated on n=128 Gaussian inputs without performance degradation.
- The same flow supplies maps in both directions between inputs and barycenter, which enables style transfer and multivariate fair regression in the experiments.
Reading between the lines
- Because Theorem 3.3 identifies h(Z) as the conditional expectation E[f(Z,S)|Z], the same conditional-variance training could plausibly handle continuously indexed families of input distributions, a direction the paper gestures at with its grid-valued condition but does not formalize beyond a finite set.
- The MNIST style-transfer result suggests the shared latent is aligned across classes; if that alignment holds, varying the weights w_s in h(z) could produce continuous morphings between distributions, which the paper does not test.
- The decreasing-weight schedule is a heuristic rather than a guaranteed optimizer; an alternative worth testing is a Lagrange-multiplier or augmented-Lagrangian scheme that enforces the pushforward constraints exactly while minimizing transport cost.
- Monitoring both the likelihood loss and the L2 cost on a benchmark with a known closed-form transport map would directly quantify the approximation gap caused by limited flow expressivity.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The manuscript proposes a primal, non-adversarial approach to computing optimal transport maps and Wasserstein-2 barycenters using conditional normalizing flows. Theorem 3.2 reformulates the p-Wasserstein distance between two absolutely continuous measures as an Lp minimization over pairs of shared-latent pushforward maps, and Theorem 3.3 shows that the weighted Wasserstein-2 barycenter of finitely many measures equals the pushforward of the conditional expectation h(z)=Σ_s w_s f(z,s), where f minimizes the expected conditional variance problem (4) subject to f(·,s)#λ=μ_s for all s. Algorithms 1 and 2 train conditional flows by maximizing likelihood while adding a decreasingly weighted L2-cost term, and the paper reports experiments on high-dimensional Gaussian and uniform benchmarks, MNIST digits, a large-number-of-inputs Gaussian set, and a multivariate fair-regression task. The appendices contain the proofs, normalizing-flow background, and full implementation details.
Significance. The theoretical reformulation is elegant and, to the best of my reading, correct: Theorem 3.3 gives a clean, intuitive characterization of Wasserstein barycenters, and the conditional architecture yields the OT maps to and from the barycenter as byproducts, avoids adversarial and bi-level training, and scales plausibly to hundreds of input distributions. The proofs in Appendix B are standard, the code is released, and the experimental section is extensive. However, the central algorithmic claim is not backed by the theory: the decreasing-weight training objective is not shown to find a minimizer of the conditional-variance problem (4), and the paper itself concedes that the flow family is not known to contain all relevant bijections. If the algorithm-to-theory gap can be closed, this would be a valuable contribution to computational OT; absent that, the paper delivers a promising heuristic with a sound but not-yet-tied theoretical foundation.
major comments (2)
- [Section 3.3.2, Algorithms 1–2 and Theorem 3.3] The training objective in Algorithms 1 and 2 is −log pθ(X,S) + ζ_t L2θ(Z,S), with ζ_t decreasing exponentially to a small value (Table 7, LOGSPACE schedules). For any parameter θ such that fθ(·,s)#λ = μ_s for all s, the likelihood term attains its global minimum, the entropy Σ_s w_s H(μ_s), regardless of which particular pushforward map is realized; the L2/conditional-variance term is the only component of the loss that selects the map for which the conclusion of Theorem 3.3 holds. Driving ζ_t toward zero therefore removes precisely the term that identifies the minimizer of the conditional-variance problem (4), and no convergence theorem for this time-varying multi-objective optimization is provided; the 'Decreasing weights' paragraph only states that the weights are chosen empirically. Consequently, even with infinite data and a sufficiently expressive flow, the final iterate is not proven to minimize (4), and h#λ is not proven to equal the barycenter; the same objection applies verbatim to the OT map x ↦ fθ(fθ^{-1}(x,1),2) produced by Algorithm 1. Since Theorem 3.3 applies only to exact minimizers of (4), this gap is load-bearing for the paper's central claim.
- [Section 3.3.1] The representational premise is also unproven, as the paper itself acknowledges: coupling-based normalizing flows are only known to be universal approximators for bijections via ill-conditioned constructions (Teshima et al., 2020; Koehler et al., 2021; Draxler et al., 2024), and well-conditioned flows are not known to approximate every bijection. Theorem 3.3 requires the specific optimal transport maps from λ to each μ_s, not just some bijection with the correct marginals, so even setting aside the optimization issue, expressivity of the chosen flow family is an additional independent limitation that prevents certifying that the algorithm's output equals the barycenter.
minor comments (8)
- [Abstract vs. Tables 3–4] The abstract's claim that the method 'compares favorably' with previous state-of-the-art methods is not uniformly supported: on the location-scatter uniform benchmark (Table 4), the proposed method's L2-UVP is worse than SCW2B at every dimension (0.998 vs 0.120 at d=2 and 5.506 vs 1.380 at d=128).
- [Section 5.2.2 and Appendix D, Tables 8–9] The text states that 'in practice all methods converge in approximately 10-20 minutes,' which contradicts Table 9, where all reported training times are on the order of 90–115 minutes; the associated claim of 'better or equal training time' is also not supported by Table 9, where the proposed method is slower than WIN on the uniform experiment (110M1S vs 90M17S at d=2).
- [Section B, proofs of Theorems 3.2 and 3.3] The proof of Theorem 3.2 refers to 'Theorem B.1' and the proof of Theorem 3.3 refers twice to 'Proposition B.1', but Appendix B contains only Lemma B.1; the labels should be made consistent.
- [Table 7] Table 7 reports only the endpoints of the LOGSPACE weight schedules (e.g., [0,−2]); since the decay rule for ζ_t is load-bearing for the behavior of Algorithms 1 and 2, the paper should state the exact schedule (number of steps, decay rate) and the final weight value for each experiment.
- [Figure 1] The three boxplot clusters are presumably the dimensions d=16, 32, 64, but the x-axis is unlabeled and the caption does not define the construction of the 'signed' relative error; please clarify the axes and the definition.
- [Section 5.2.5] 'We fail to reject the null hypothesis of no association' is the desired outcome for the fairness experiment, but with p-values of 0.80–0.99 and no reported sample size, test statistic, or confidence intervals, the conclusion of 'almost perfect uncorrelatedness' is stronger than the evidence presented.
- [Section 2.1 and Section 4] There are several small inconsistencies: the statement labeled 'Lemma 3.1' in Section 3.2 is referred to as 'Theorem 3.1' in Section 2.1; in Section 4, 'paramatrize' should be 'parametrize' and 'loose accuracy' should be 'lose accuracy'; and the reference to Atanov et al. (2019) contains '920' instead of the publication year.
- [Section 5.2.4, Table 5] The text states that the results show 'no degradation in performance as the number of input distributions increases,' but Table 5 shows L2-UVP rising from 0.057 at n=4 to 0.119 at n=16, roughly a two-fold worsening; the statement should be qualified.
Circularity Check
No significant circularity: the theoretical results rest on external theorems and standard measure theory, and the experimental validation uses independent baselines and known closed-form OT maps.
full rationale
The paper's derivation chain is self-contained and does not reduce its predictions to its inputs. Theorem 3.2 (the constrained L^p(λ) reformulation of the OT problem) is proved from Brenier's theorem via Gangbo-McCann, an external mathematical result, and the proof is given in Appendix B. Theorem 3.3 (the conditional-variance characterization of Wasserstein-2 barycenters) is proved using standard conditional-expectation identities and the external existence/uniqueness result for barycenters (Agueh-Carlier; Brizzi et al. 2025). The barycenter is not defined as the output of the fitted flow; it is characterized independently and then shown to equal h#λ for any minimizer of problem (4). Algorithm 2's training objective is heuristic: it combines negative log-likelihood with a decreasingly weighted L2 cost, and the paper provides no convergence theorem showing the final iterate minimizes (4). This is a genuine rigor/limitation concern, but it is not circularity under the definitions used here: the output map and barycenter are not defined in terms of the fitted loss, and the method's correctness is not assumed by its own construction. The numerical experiments compare against external baselines (SCW2B, WIN) and against known optimal transport maps and barycenters obtained from a fixed-point algorithm and closed-form Gaussian/Uniform solutions. The location-scatter validation uses the same affine family that motivates the architecture, but the reference quantities are computed independently, so this does not make the evaluation circular. There are no load-bearing self-citations by the present authors; the cited flow-expressivity caveat (Draxler et al.) is an external limitation, not a self-referential guarantee. Overall, no step in the claimed derivation is equivalent to its own input by construction, and no fitted parameter is renamed as a prediction.
Assumptions & free parameters
free parameters (1)
- L2-cost weight schedule (ζ_t) =
LOGSPACE [0,-2] or [1,-2] depending on experiment (Table 7)
assumptions (5)
- standard math Brenier's theorem for the existence and uniqueness of optimal transport maps between absolutely continuous measures (Gangbo & McCann 1996).
- standard math The Wasserstein-2 barycenter of finitely many measures in P_{2,ac}(R^d) exists, is unique, and is absolutely continuous (Agueh & Carlier 2011; Brizzi et al. 2025).
- standard math Standard measure-theoretic tools: change-of-variables formula, uniqueness of conditional expectation, and properties of L^p spaces.
- domain assumption The input distributions μ_s belong to P_{2,ac}(R^d), i.e. they are absolutely continuous with finite second moments.
- ad hoc to paper The conditional normalizing flow family can represent the relevant optimal transport maps, and the training procedure (maximum likelihood plus decreasing L2 weight) converges to a joint minimizer.
Cite this review
Pith. "Pith review of Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows." pith.science (2026). https://pith.science/paper/JXB7BMDV
@misc{pith2026250522364,
author = {Pith},
title = {Pith review of: Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows},
year = {2026},
howpublished = {\url{https://pith.science/paper/JXB7BMDV}},
note = {Machine review of arXiv:2505.22364}
}
read the original abstract
We present a novel method for efficiently computing optimal transport maps and Wasserstein barycenters in high-dimensional spaces. Our approach uses conditional normalizing flows to approximate the input distributions as invertible pushforward transformations from a common latent space. This makes it possible to directly solve the primal problem using gradient-based minimization of the transport cost, unlike previous methods that rely on dual formulations and complex adversarial optimization. We show how this approach can be extended to compute Wasserstein barycenters by solving a conditional variance minimization problem. A key advantage of our conditional architecture is that it enables the computation of barycenters for hundreds of input distributions, which was computationally infeasible with previous methods. Our numerical experiments illustrate that our approach yields accurate results across various high-dimensional tasks and compares favorably with previous state-of-the-art methods.
Figures
Figures from the paper (4 more)
Forward citations
Cited by 1 Pith paper
-
Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation
A mini-batch Wasserstein gradient-flow algorithm computes scalable and label-aware Wasserstein barycenters, with empirical gains on domain adaptation.
Reference graph
Works this paper leans on
-
[5]
Density estima- tion using Real NVP.arXiv preprint arXiv:1605.08803,
Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density estima- tion using Real NVP.arXiv preprint arXiv:1605.08803,
-
[7]
Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks
Fan, J., Taghvaei, A., and Chen, Y . Scalable computations of Wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.04462,
work page Pith review arXiv 2007
-
[8]
L., Loubes, J.-M., and Rigollet, P
Gouic, T. L., Loubes, J.-M., and Rigollet, P. Projec- tion to fairness in statistical learning.arXiv preprint arXiv:2005.11720,
arXiv 2005
-
[10]
M., Filippov, A., and Burnaev, E
Korotin, A., Li, L., Genevay, A., Solomon, J. M., Filippov, A., and Burnaev, E. Do neural optimal transport solvers work? A continuous Wasserstein-2 benchmark.Advances in Neural Information Processing Systems, 34:14593– 14605, 2021a. Korotin, A., Li, L., Solomon, J., and Burnaev, E. Continuous Wasserstein-2 barycenter estimation without minimax optimizati...
-
[12]
Lu, G., Zhou, Z., Shen, J., Chen, C., Zhang, W., and Yu, Y . Large-scale optimal transport via adversarial training with cycle-consistency.arXiv preprint arXiv:2003.06635,
arXiv 2003
-
[15]
The intrinsic dimension of images and its impact on learning.arXiv preprint arXiv:2104.08894,
Pope, P., Zhu, C., Abdelkader, A., Goldblum, M., and Gold- stein, T. The intrinsic dimension of images and its impact on learning.arXiv preprint arXiv:2104.08894,
-
[17]
12 Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows A. Auxiliary Results The following are some standard results from measure and integration theory that are used in this paper. Lemma A.1(Change of variables).Let f:X→Y be a measurable function between two measurable spaces (X,A) and (Y,B) . If µ is probabili...
work page 2007
-
[18]
Proof. R Rd ∥f(x)∥µ(dx) = 0 if and only if ∥f(x)∥= 0µ -almost surely (see, for instance, Proposition 2.16 in Folland (1999), which is the case if and only iff(x)is the zero vectorµ-almost surely. Lemma A.3.Letµ, ν∈ Pp(Rd)for somep≥1. Iff∈B(µ, ν)and ˜fis a(µ, ν)-inverse off, then (i)∥f∥ ∈L p(µ), (ii) ˜fis aν-a.s. unique(µ, ν)-inverse off, (iii) ˜f∈B(ν, µ),...
work page 1999
Show all 18 references
-
[1781]
Wasserstein style transfer.arXiv preprint arXiv:1905.12828,
Mroueh, Y . Wasserstein style transfer.arXiv preprint arXiv:1905.12828,
1905 arXiv
-
[1982]
On the universality of volume-preserving and coupling-based normalizing flows.arXiv preprint arXiv:2402.06578,
Draxler, F., Wahl, S., Schn ¨orr, C., and K ¨othe, U. On the universality of volume-preserving and coupling-based normalizing flows.arXiv preprint arXiv:2402.06578,
-
[2002]
Shen, J., Qu, Y ., Zhang, W., and Yu, Y
DOI: https://doi.org/10.24432/C53W3X. Shen, J., Qu, Y ., Zhang, W., and Yu, Y . Wasserstein distance guided representation learning for domain adaptation. In Proceedings of the AAAI Conference on Artificial Intelli- gence, pp. 4058–4065,
-
[2014]
D., Bui, H., and Phung, D
Dam, N., Hoang, Q., Le, T., Nguyen, T. D., Bui, H., and Phung, D. Three-player Wasserstein GAN via amortised duality. InInternational Joint Conference on Artificial Intelligence 2019, pp. 2202–2208. Association for the Advancement of Artificial Intelligence (AAAI),
2019
-
[2017]
Semi-conditional normalizing flows for semi- supervised learning.arXiv preprint arXiv:1905.00505, 920,
Atanov, A., V olokhova, A., Ashukha, A., Sosnovik, I., and Vetrov, D. Semi-conditional normalizing flows for semi- supervised learning.arXiv preprint arXiv:1905.00505, 920,
1905 arXiv
-
[2019]
NICE: Non-linear Independent Components Estimation.arXiv preprint arXiv:1410.8516,
Dinh, L., Krueger, D., and Bengio, Y . NICE: Non-linear Independent Components Estimation.arXiv preprint arXiv:1410.8516,
-
[2021]
On the regularization of Wasserstein GANs.arXiv preprint arXiv:1709.08894,
Petzka, H., Fischer, A., and Lukovnicov, D. On the regularization of Wasserstein GANs.arXiv preprint arXiv:1709.08894,
-
[2023]
Adversarial computation of optimal trans- port maps.arXiv preprint arXiv:1906.09691,
Leygonie, J., She, J., Almahairi, A., Rajeswar, S., and Courville, A. Adversarial computation of optimal trans- port maps.arXiv preprint arXiv:1906.09691,
1906 arXiv
-
[2024]
Wasserstein-2 generative networks.arXiv preprint arXiv:1909.13082,
Korotin, A., Egiazarian, V ., Asadulaev, A., Safin, A., and Burnaev, E. Wasserstein-2 generative networks.arXiv preprint arXiv:1909.13082,
1909 arXiv
-
[2025]
C., Caterini, A
Brown, B. C., Caterini, A. L., Ross, B. L., Cresswell, J. C., and Loaiza-Ganem, G. Verifying the union of manifolds hypothesis for image data.arXiv preprint arXiv:2207.02862,
Reviewed August 7, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.