REVIEW 4 major objections 5 minor 1 cited by
Optimal transport mapping via input convex neural networks
T0 review · 4 major / 5 minor · reviewed 2026-08-14 · deepseek-v4-flash
Pith's one-line read A minimax game over convex functions recovers the exact optimal transport map.
desk verdict A genuinely new minimax formulation for learning the 2-Wasserstein map that is worth engaging, but the theory doesn't cover the practical algorithm and the initialization-independence claim is demonstrated only qualitatively. 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 carrying object is the functional $V_{P,Q}(f,g) = -\mathbb{E}_P[f(X)] - \mathbb{E}_Q[\langle Y, \nabla g(Y)\rangle - f(\nabla g(Y))]$ over pairs of convex functions, together with the identity $\langle y, \nabla g(y)\rangle - f(\nabla g(y)) \le f^*(y)$, equality holding when $g = f^*$. This converts the hard Kantorovich feasibility constraint into an unconstrained saddle problem: the inner minimization selects $g$ as the convex conjugate of $f$, and the outer maximization selects the Brenier potential whose gradient is the optimal map. Input convex neural networks, networks whose output is convex in the input, supply the trainable parameterization of the convex functions, and a universality result for these networks is used to justify replacing the full convex class by the neural class.
What would settle it
Take $Q = \mathcal{N}(0,I_d)$ and $P = \mathcal{N}(\mu,I_d)$, where the exact optimal map is $T^*(y) = y + \mu$, train the proposed algorithm from many different initializations, and compare each learned $\nabla g$ with $T^*$ in $L^2(Q)$ while tracking the converged minimax value against the known squared Wasserstein distance $\|\mu\|^2/2$; if the maps vary with initialization or the objective converges to a different value as network width grows, the saddle-point-to-map claim is disproved.
Extended reading notes
Core claim
The central discovery is Theorem 3.3: for distributions $P$ and $Q$ with finite second moments, with $Q$ having a density, the squared 2-Wasserstein distance equals $\sup_{f \in \mathrm{CVX}(P), f^* \in L^1(Q)} \inf_{g \in \mathrm{CVX}(Q)} V_{P,Q}(f,g) + C_{P,Q}$, where $V_{P,Q}(f,g) = -\mathbb{E}_P[f(X)] - \mathbb{E}_Q[\langle Y, \nabla g(Y)\rangle - f(\nabla g(Y))]$, and there is an optimal pair $(f_0,g_0)$ with $\nabla g_0$ equal to the optimal transport map from $Q$ to $P$. The proof rests on the inequality $\langle y, \nabla g(y)\rangle - f(\nabla g(y)) \le f^*(y)$, with equality at $g = f^*$, together with the Knott-Smith criterion and Brenier's theorem. The paper further claims that the proposed algorithm, which parameterizes $f$ and $g$ by input convex neural networks and alternates gradient updates, learns this optimal map independent of initialization, and that using gradients of convex networks permits discontinuous transport maps with sharp boundaries.
Load-bearing premise
The result rests on the assumption that the infinite-dimensional minimax over all convex functions is attained by some pair $(f_0,g_0)$ and that $Q$ has a density; in the practical version, the stability bound additionally assumes $f$ is $\alpha$-strongly convex, which the ICNN parameterization does not guarantee.
Editorial extensions
If this is right
- An exact solution of the minimax gives the exact optimal transport map, not merely an arbitrary feasible transport.
- Deep generative models trained this way inherit uniqueness: the same generator emerges from different initializations, so the model is not one of many arbitrary push-forwards.
- Because the gradient of a convex ReLU-based network can have kinks, the learned map can jump across gaps, producing sharp boundaries for disconnected target supports.
- Near-optimal solutions satisfy an explicit $L^2$ error bound proportional to the objective gaps divided by the strong-convexity parameter $\alpha$, linking optimization error to transport error.
- The framework turns optimal transport theory into a practical training signal for generative modeling, with robustness and discontinuity as concrete advantages over adversarial training.
Reading between the lines
- The stability bound requires $f$ to be $\alpha$-strongly convex, a property the ICNN parameterization does not enforce; adding a small quadratic regularizer to $f$ would make the bound applicable, a modification the paper does not explore.
- The minimax identity is representation-agnostic: any sufficiently expressive convex function class could replace ICNNs, and the same consistency would hold up to an approximation error that could be bounded by combining the universality result with Theorem 3.6.
- The finite-sample behavior is left open; deriving sample complexity for the empirical minimax objective is a natural extension, and the paper's stability result only partially addresses it.
- The dependence on $Q$ having a density suggests the method will need modification when the source distribution has atoms, since the differentiability-almost-everywhere argument in the proof would break.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes a new approach for learning optimal transport maps under the squared-Euclidean (W2) cost from samples. The authors introduce a minimax formulation over pairs of convex functions, derived from the Kantorovich dual via the convexification trick, and show in Theorem 3.3 that the solution of this minimax problem yields the optimal transport map as the gradient of one of the convex potentials. They parameterize the convex functions with input convex neural networks (ICNNs) and propose Algorithm 1, a stochastic minimax training procedure with a soft-convexity penalty for the second network. Numerical experiments on 2D synthetic distributions, high-dimensional Gaussians, Gaussian mixtures, and MNIST latents demonstrate the method's ability to recover optimal maps, including discontinuous maps, and its relative robustness to initialization compared to WGAN-type baselines.
Significance. The paper makes a useful methodological contribution by providing a clear, theory-grounded objective for learning W2-optimal transport maps, and demonstrates empirically that representing the map as the gradient of an ICNN allows sharp separation of disconnected target supports—an advantage over continuous network maps. The theoretical results, if fully established, would provide consistency and a stability estimate linking optimization gaps to map error. The experimental appendix is detailed and supports reproducibility. However, as presented, the proof of Theorem 3.3 omits the existence argument for the outer maximizer, and the stability theorem's proof contains an error in the constant; additionally, the bridge from the infinite-dimensional theory to the ICNN/penalized algorithm is not quantified. These gaps are fixable but require careful revision.
major comments (4)
- [Section 3 (Theorem 3.3) and Appendix A] The proof sketch and Appendix A establish that for a fixed f, the inner supremum over g in CVX(Q) is attained by g=f^*, but they do not establish the existence of an f0 attaining the outer supremum in CVX(P) with f0^∗ ∈ L1(Q). The statement that '(f0,g0) achieves the infimum and supremum' is invoked without proof; this is a standard existence result in optimal transport, but the paper should either supply the argument or cite the precise theorem (e.g., Villani 2003, Theorem 2.9 or 2.12) in the proof. Without this, the saddle-point characterization and the claim that ∇g0 is the optimal transport map are not fully justified.
- [Section 3.2 (Theorem 3.6) and Appendix B] The proof uses the triangle inequality on the L2-norm after deriving ||∇g−∇f^*||^2_{L2(Q)} ≤ (2/α)ϵ1 and ||∇f^*−∇g0||^2_{L2(Q)} ≤ (2/α)ϵ2. For any two square-integrable functions u and v, ||u+v||^2 ≤ 2||u||^2 + 2||v||^2, so the triangle inequality yields ||∇g−∇g0||^2_{L2(Q)} ≤ (4/α)(ϵ1+ϵ2), not (2/α)(ϵ1+ϵ2) as stated. The constant in Theorem 3.6 thus appears to be incorrect under the proof as written; the authors should either correct the theorem to a constant of 4/α or find an alternative argument to retain the stated bound.
- [Section 3.1 and Algorithm 1] The theoretical guarantees in Theorems 3.3 and 3.6 are stated for exact optimization over the infinite-dimensional spaces CVX(P) and CVX(Q). The practical algorithm, however, optimizes over ICNN parameterizations and relaxes the convexity constraint on g by replacing it with the soft penalty R(θg) in Eq. (8). The paper does not provide a consistency result or an approximation bound that accounts for these parametric and penalization errors. Consequently, the abstract's claim that the learned mapping is 'optimal independent of how we initialize the neural networks' is not directly supported by the proven theorems; it is an empirical observation that appears consistent with the experiments but requires either a gap analysis or a more careful statement of the claim.
- [Section 3.2 (Theorem 3.6)] The stability bound in Theorem 3.6 assumes f is α-strongly convex, but the ICNN parameterization in Section 3.1 only enforces convexity (with nonnegative weights), and no strong-convexity regularization or verification is described in the experiments of Section 4. As stated, the theorem is vacuous if no learned potential satisfies the assumption; the authors should either discuss how α can be enforced or estimated for ICNNs or explicitly label the result as conditional on an unverified hypothesis.
minor comments (5)
- [Section 3, Eq. (4)] The expression 'inf_{f∈CVX(P)} E_P[f(X)] + E_Q[f^*(Y)]' should be parenthesized to make clear that the infimum applies to the sum of the two expectations, not only to the first term.
- [Table 1 caption] The caption states that the source and target are 728-dimensional Gaussians, whereas the text in Section 4.4 specifies d=784; this is a typo.
- [Section 4.4, Table 1] The reported metric ||µT(Q)-µ||2 in the first row of Table 1 is scale-dependent; the relative error in the second row is more informative, and the first row could be complemented with a normalized version or a discussion of this dependence.
- [Algorithm 1] The step labeled 'Projection' is a thresholding of weights to nonnegative values; using a more descriptive name such as 'Weight clipping' would avoid confusion with projection onto a feasibility set.
- [Appendix A] In the proof of Theorem 3.3, the differentiability of f^* Q-a.e. is asserted to follow by arguments similar to those for g, but the argument is not spelled out; adding a sentence that f^* is convex and finite Q-a.e. (by f^*∈L1(Q)) would close this gap.
Circularity Check
No significant circularity: the minimax formulation is derived from classical optimal transport theorems, and no fitted quantity is renamed as a prediction.
full rationale
The paper's central result, Theorem 3.3, is derived from classical tools: the Kantorovich duality (Villani, Theorem 1.3), the semi-dual formulation from Villani (Theorem 2.9), and the Knott-Smith/Brenier characterization of optimal maps (Theorem 3.1). The proof of Theorem 3.3 in Appendix A reduces the supremum over convex g to E_Q[f^*(Y)] using the pointwise Fenchel-Young inequality and a.e. differentiability of convex functions; this is a standard convex-analytic identity, not an assumption that includes the target result. The estimated transport map is the gradient of the inner convex function, and the claim that the saddle point yields the optimal map is a direct corollary of Brenier's theorem rather than a consequence of the paper's own parameterization. The self-citation to Taghvaei & Jalali (2019) is contextual: it credits prior work for using ICNNs in Wasserstein approximation, but the minimax formulation, its proof, and the stability analysis do not rely on the cited paper's theorems. The Gaussian experiments compare the trained objective to the exact W2 value, but they also report an independent out-of-sample transport accuracy metric (mean of transported distribution versus target mean), so the evaluation is not merely a restatement of the training objective. Concerns about attainment of the supremum/infimum in infinite-dimensional spaces, the gap between ICNN convexity and the theorem's hypotheses, and the unverified strong-convexity assumption in Theorem 3.6 are correctness and rigor concerns, not circularity: the paper does not define its target result into its assumptions or use a self-citation to force that result. The derivation chain is therefore self-contained with respect to its classical inputs.
Assumptions & free parameters
free parameters (2)
- Regularization constant λ for g network weights =
1.0 (checkerboard), 0.1 (Gaussian-to-Gaussian), 0.01 (Gaussian-to-mixture), 1.0 (MNIST experiments)
- Hyperparameters: hidden size, layers, batch size, generator iterations K, learning rate schedule =
Listed per experiment in Appendix C
assumptions (5)
- domain assumption Q admits a density with respect to Lebesgue measure on Rd
- domain assumption The convex function spaces CVX(P) and CVX(Q) contain a minimizer/maximizer pair (f0, g0) for the minimax problem
- domain assumption Finite second-order moments for P and Q
- domain assumption The ICNN class can approximate any convex function over the relevant domain
- ad hoc to paper f is α-strongly convex in Theorem 3.6
Cite this review
Pith. "Pith review of Optimal transport mapping via input convex neural networks." pith.science (2026). https://pith.science/paper/JXU6T446
@misc{pith2026190810962,
author = {Pith},
title = {Pith review of: Optimal transport mapping via input convex neural networks},
year = {2026},
howpublished = {\url{https://pith.science/paper/JXU6T446}},
note = {Machine review of arXiv:1908.10962}
}
read the original abstract
In this paper, we present a novel and principled approach to learn the optimal transport between two distributions, from samples. Guided by the optimal transport theory, we learn the optimal Kantorovich potential which induces the optimal transport map. This involves learning two convex functions, by solving a novel minimax optimization. Building upon recent advances in the field of input convex neural networks, we propose a new framework where the gradient of one convex function represents the optimal transport mapping. Numerical experiments confirm that we learn the optimal transport mapping. This approach ensures that the transport mapping we find is optimal independent of how we initialize the neural networks. Further, target distributions from a discontinuous support can be easily captured, as gradient of a convex function naturally models a {\em discontinuous} transport mapping.
Figures
Figures from the paper (3 more)
Forward citations
Cited by 1 Pith paper
-
Input convex neural networks as surrogates in mathematical optimisation
ICNN surrogates admit an LP-only embedding whose epigraph plus concave envelope equals the convex hull of the network's graph, enabling a branch-and-bound solver that avoids binary variables.
Reference graph
Works this paper leans on
-
[1]
Amos, B., Xu, L., and Kolter, J. Z. Input convex neural networks. arXiv preprint arXiv:1609.07152,
-
[2]
andP is the law of random variable Y , where Y = X +Z with X∼ Unif({(1, 0), ( 1√ 2, 1√ 2)}, (0, 1), (−1√ 2, 1√ 2), (−1, 0), (−1√ 2,−1√ 2), (0,−1), ( 1√ 2,−1√ 2)}) andZ∼N (0, 0.5I2). Architecture details. For our Algorithm 1, we parametrize both the convex functionsf andg by ICNNs. Both these ICNN networks have equal number of nodes for all the hidden laye...
work page 2019
-
[9]
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114,
-
[12]
Adversarial computation of optimal trans- port maps
Leygonie, J., She, J., Almahairi, A., Rajeswar, S., and Courville, A. Adversarial computation of optimal trans- port maps. arXiv preprint arXiv:1906.09691,
arXiv 1906
-
[14]
On the regularization of wasserstein gans
Petzka, H., Fischer, A., and Lukovnicov, D. On the regularization of wasserstein gans. arXiv preprint arXiv:1709.08894,
-
[15]
Uncoupled isotonic regression via minimum Wasserstein deconvolution
Rigollet, P. and Weed, J. Uncoupled isotonic regression via minimum wasserstein deconvolution. arXiv preprint arXiv:1806.10648,
-
[16]
B., Flamary, R., Courty, N., Rolet, A., and Blondel, M
Seguy, V ., Damodaran, B. B., Flamary, R., Courty, N., Rolet, A., and Blondel, M. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv:1711.02283,
-
[19]
On Scalable and Efficient Computation of Large Scale Optimal Transport
Xie, Y ., Chen, M., Jiang, H., Zhao, T., and Zha, H. On scalable and efficient computation of large scale optimal transport. arXiv preprint arXiv:1905.00158,
work page Pith review arXiv 1905
Show all 23 references
-
[20]
and Karniadakis, G
Yang, L. and Karniadakis, G. E. Potential flow genera- tor with l_2 optimal transport regularity for generative models. arXiv preprint arXiv:1908.11462,
1908 arXiv
-
[23]
propose to learn the functionf in a semi-discrete setting, where one of the marginals is assumed to be a discrete distribution supported on a set ofN points{y1,...,y N}⊂ Rd, and the other marginal is assumed to have a continuous density with compact convex support Ω⊂ Rd. They ...
2019
-
[64]
D. Further discussion of related work The idea of solving the semi-dual optimization problem (4) is classically considered in (Chartrand et al., 2009), where the authors derive a formula for the functional derivative of the objective function with respect tof and propose to so...
2009
-
[1781]
and Cuturi, M
Muzellec, B. and Cuturi, M. Subspace detours: Building transport plans that are optimal on subspace projections. arXiv preprint arXiv:1905.10099,
1905 arXiv
-
[1998]
Lei, N., Su, K., Cui, L., Yau, S.-T., and Gu, D. X. A geometric view of optimal transportation and generative model. arXiv preprint arXiv:1710.05488,
-
[2000]
Smooth and sparse op- timal transport
Blondel, M., Seguy, V ., and Rolet, A. Smooth and sparse op- timal transport. arXiv preprint arXiv:1710.06276,
-
[2003]
Improving the improved training of wasserstein gans: A consistency term and its dual effect
Wei, X., Gong, B., Liu, Z., Lu, W., and Wang, L. Improving the improved training of wasserstein gans: A consistency term and its dual effect. arXiv preprint arXiv:1803.01541,
-
[2009]
Optimal control via neural networks: A convex approach
Chen, Y ., Shi, Y ., and Zhang, B. Optimal control via neural networks: A convex approach. arXiv preprint arXiv:1805.11835,
-
[2013]
Wasserstein-2 generative networks
Korotin, A., Egiazarian, V ., Asadulaev, A., and Burnaev, E. Wasserstein-2 generative networks. arXiv preprint arXiv:1909.13082,
1909 arXiv
-
[2014]
Domain adaptation for object recognition: An unsupervised approach
Gopalan, R., Li, R., and Chellappa, R. Domain adaptation for object recognition: An unsupervised approach. In 2011 international conference on computer vision , pp. 999–1006. IEEE,
2011
-
[2015]
and Jalali, A
Taghvaei, A. and Jalali, A. 2-wasserstein approximation via restricted convex potentials with application to improved training for gans. arXiv preprint arXiv:1902.07197,
1902 arXiv
-
[2016]
Wasserstein gan
Arjovsky, M., Chintala, S., and Bottou, L. Wasserstein gan. arXiv preprint arXiv:1701.07875,
-
[2017]
Mode collapse and regularity of optimal transportation maps
Guo, Y ., An, D., Qi, X., Luo, Z., Yau, S.-T., Gu, X., et al. Mode collapse and regularity of optimal transportation maps. arXiv preprint arXiv:1902.02934,
1902 arXiv
-
[2018]
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980,
-
[2019]
and Rigollet, P
Hütter, J.-C. and Rigollet, P. Minimax rates of estima- tion for smooth optimal transport maps. arXiv preprint arXiv:1905.05828,
1905 arXiv
Reviewed August 14, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.