REVIEW 4 major objections 6 minor 35 references
A generalized Wasserstein-2 distance approach for efficient reconstruction of random field models using stochastic neural networks
T0 review · 4 major / 6 minor · reviewed 2026-08-06 · deepseek-v4-flash
Pith's one-line read The paper proves that stochastic neural networks can approximate any sufficiently regular random field with mixed continuous and categorical outputs to arbitrary accuracy under a generalized Wasserstein-2 distance, and trains them with a…
desk verdict A plausible extension of local W2 training to mixed continuous/categorical random variables, but the central approximation theorem depends on an unstated prior theorem and hand-waved regularity conditions; worth a serious referee who expects revision. 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 load-bearing object is a modified norm on the mixed vector, $\|y\|^2=\lambda\sum_{i=1}^{d_1} y_i^2 + \sum_{j=d_1+1}^d \hat\delta_{y_j,0}$, where $\hat\delta_{y_j,0}=4y_j^2$ for $|y_j|\le 1/2$ and $1$ otherwise. This makes categorical mismatches contribute a bounded, almost quadratic penalty instead of an unbounded Euclidean distance, and it defines the generalized Wasserstein-2 distance $\hat W_2$ between conditional distributions. The approximation proof smooths the categorical components with a mollifier so that the known SNN universal approximation theorem for continuous random fields applies, then transfers the resulting bound back to the original mixed measure by triangle inequalities. The training side uses a local empirical version of $\hat W_2^2$, with a differentiable surrogate for the categorical discrepancy, quadratic when prediction is near the label and cosine-based when it is far, so that gradient descent can train through discrete targets.
What would settle it
Take a mixed random field with bounded outcomes whose categorical class probabilities jump discontinuously at a threshold in $x$, so the uniform Lipschitz condition in Eq. (2.8) fails, and train an SNN with the proposed local squared loss; if the learned distributions still converge to the target under $\hat W_2$, the Lipschitz assumption is unnecessary, and if they do not, the theorem's stated scope is confirmed. A complementary check is to verify the technical regularity conditions on $D$ named in [34, Appendix H] for a concrete domain used in the experiments; if those conditions are violated, the proof of Theorem 2.1 is incomplete as written.
Extended reading notes
Core claim
On its own terms, the paper's central discovery is Theorem 2.1: for any random field $y_x=y(x;\omega)$ whose conditional distributions $f_x$ are uniformly bounded, uniformly Lipschitz in $x$ under the generalized Wasserstein-2 metric, and satisfy the smoothness conditions in Assumption 2.1, and for any $\epsilon_1>0$, there exists a stochastic neural network whose output distribution $\hat f_x$ satisfies $\int_D \hat W_2^2(f_x,\hat f_x)\,\nu(dx)\le \epsilon_1$. Continuous outputs are rounded at test time to deliver categorical predictions, while training uses a surrogate of the categorical discrepancy that keeps gradients nonzero. The paper also proves Theorem 2.2, a finite-sample generalization bound showing that the local empirical loss approaches the true squared distance at a rate controlled by $1/\sqrt N$, the neighborhood size $\delta$, and a dimension-dependent empirical-measure term. These two results underwrite the practical recipe: minimize the generalized local squared Wasserstein-2 loss over the means and variances of the stochastic network weights.
Load-bearing premise
The load-bearing premise is the regularity package in Assumption 2.1, above all the uniform Lipschitz continuity of the true conditional distributions in $x$ and the technical conditions on the domain $D$ that the proof inherits from an earlier continuous-field approximation theorem and simply assumes hold, so that the black-box universal approximation result can be applied.
Editorial extensions
If this is right
- If Theorem 2.1 is correct, SNNs are a universal model class for mixed-output random fields: no separate continuous and categorical architectures are needed for distributional reconstruction.
- The local squared loss provides a finite-data objective whose error relative to the true distance is controlled by three explicit terms, sample size, local neighborhood statistics, and Lipschitz regularity, so users can predict when reconstruction will be accurate.
- Because the loss is differentiable and local, it can be used as a drop-in replacement for mean squared error or cross-entropy in uncertainty-quantification pipelines for mixed targets.
- The learned per-input randomness comes from sampling stochastic weights, so repeated forward passes at a fixed $x$ give a Monte Carlo estimate of the full conditional distribution.
- For multivariate categorical targets, encoding them as a single categorical variable improves both accuracy and runtime, matching the dimension dependence in the bound.
Reading between the lines
- The categorical penalty saturates at 1 for well-separated labels, so the loss effectively ignores how different two wrong categories are; permuting the category codes should leave training unchanged, a testable consequence the paper does not run.
- The mollification step is a general transfer principle: any universal approximation theorem for continuous random fields should lift to mixed random fields under the same assumptions, so the proof strategy could be reused for other stochastic network architectures without redoing the argument.
- The bound in Theorem 2.2 has three competing terms depending on the neighborhood size $\delta$; balancing $1/\sqrt N$, $h(N(x,\delta),d)$, and $L\delta$ would give a principled rule for choosing $\delta$, which the paper leaves to heuristic search.
- The method only needs local empirical distributions of outputs, so it should extend to online or streaming settings where minibatches are refreshed during training, though the paper does not analyze this regime.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes a generalized Wasserstein-2 distance for random field models whose outputs contain both continuous and categorical components. It defines a generalized squared W2 loss, proves a universal approximation theorem for stochastic neural networks under this distance (Theorem 2.1) and a generalization error bound for the empirical local loss (Theorem 2.2), then introduces a differentiable surrogate loss and tests the method on classification, mixed continuous/categorical regression (abalone), and a gene-regulatory ODE/Markov jump system.
Significance. If the theoretical results are made fully rigorous, the paper offers a useful principled extension of local squared Wasserstein training to categorical and mixed outputs, a setting where the usual W2 distance is not directly applicable. The differentiable surrogate with a detach-based rounding trick is a practical contribution, and the three diverse UQ experiments plus the sensitivity analyses give useful empirical evidence. The main limitation is that the central approximation theorem currently rests on an unstated external theorem and a smoothing construction that is not fully justified, so the theoretical support does not yet match the strength of the claims.
major comments (4)
- [Appendix A, Eq. (A.9)] The proof of Theorem 2.1 reduces the approximation of the smoothed measure f_{ε,x} to the universal approximation theorem of [34, Appendix H], but that theorem is neither stated nor proved in this manuscript, and the text explicitly says 'For simplicity, we assume those conditions hold here' concerning technical regularity conditions on D. Since [34] is a preprint and its hypotheses are not checkable from the present paper, Theorem 2.1 is conditional at its load-bearing step. Please state the external theorem in full and either prove it or verify all of its hypotheses, or replace the reduction with a self-contained proof.
- [Appendix A, Eqs. (A.2)-(A.4)] The definition of f_{ε,x} in Eq. (A.2) is not a probability density as written: it appears to depend on a single discrete y with no summation over the categorical support, and it is not normalized. Moreover, the convolution with φ_ε acts only on the categorical coordinates, so the claimed smoothness of f_{ε,x} in all d coordinates does not follow from the uniform continuity of f_x in the continuous components stated in Assumption 2.1.7. This affects the verification of Eqs. (A.4) and (A.9). Please define f_{ε,x} as a genuine convolution in all coordinates (for example, summing over y_2 and mollifying the continuous coordinates as well) and prove that the mixed-derivative bounds in Assumption 2.1 are inherited.
- [Section 2, Eqs. (2.1)-(2.2); Appendix B, Eq. (B.7)] The proof of Theorem 2.2 applies the triangle inequality to \hat W_2, but the cost defined in Eq. (2.1) is not shown to make \hat W_2 a metric and it is not a norm: for a one-dimensional continuous component, \|(0.8,0,\dots,0)\| = 1, while 2\|(0.4,0,\dots,0)\| = 1.6, so homogeneity fails. The paper calls Eq. (2.1) a norm and claims the coefficient 4 ensures this, but that claim is false. A lemma establishing that the generalized distance satisfies the triangle inequality (or a different argument) is needed before Eq. (B.7) can be used.
- [Section 2.3, Eqs. (2.23)-(2.26)] The loss actually minimized in all experiments is the differentiable surrogate |\cdot|_1 with the detach/straight-through estimator in Eq. (2.25), while Theorems 2.1 and 2.2 are stated for the generalized \hat W_2 built from Eq. (2.1). The equality between |y-\hat y|_1 and \|y-\hat y\| holds only when both variables are rounded to integers, which is not the training phase. No argument shows that minimizing the surrogate controls \hat W_2 or that the detach operation preserves the objective. Please either prove a relation between the surrogate and \hat W_2, or explicitly present Eq. (3.1) as a heuristic whose validation is empirical.
minor comments (6)
- [Section 2, Eq. (2.2)] The notation \hat\delta_{y_j,0} has one argument, but Eqs. (2.14)-(2.15) and (D.1) use a two-argument form \hat\delta_{y_i,\hat y_i}; please define this two-argument version explicitly.
- [Section 2, Eq. (2.1)] The notation \|\cdot\|_2 is used both for the generalized quantity in Eq. (2.1) and for the Euclidean ℓ2 norm in Eq. (2.7), which is confusing; please use distinct symbols.
- [Section 3, Examples 3.1-3.4] Assumptions 2.1.3 and 2.1.4 require the conditional distributions to be uniformly Lipschitz in x under the generalized W2 distance, but the experiments provide no diagnostic or discussion of whether this regularity condition is plausible for the tested models.
- [Section 3, Figures 2-4 and Table 2] All experimental results appear to be single-run and no error bars, confidence intervals, or multiple-seed variability are reported, which makes it difficult to judge whether the observed differences between methods are significant.
- [Section 3.2, paragraph after Eq. (3.3)] The sentence 'For predicting the categorical sex variable on the testing set' appears to be a copy-paste from Example 3.3; in Example 3.2 the target is a multilabel binary vector, not a sex variable.
- [Section 2.3, Eq. (2.25)] The detach operation in round_1 is a straight-through estimator and deserves a more explicit explanation, since it is central to the practical training procedure but is mentioned only in one sentence.
Circularity Check
Theorem 2.1's approximation guarantee is inherited from the authors' own prior preprint [34, Appendix H], whose hypotheses are not verified and whose conditions on D are assumed; the central claim therefore reduces to a self-citation chain.
-
self citation load bearing
[Appendix A, proof of Theorem 2.1, after Eq. (A.9)]
"Combining Eqs. (A.4), (A.5), the uniform continuity of fε,x(z), and the Lipschitz condition Eq. (A.8), the assumptions in the universal approximation ability theorem of SNNs in [34, Appendix H] hold. Therefore, for any ε0 > 0, from [34, Appendix H], there exists an SNN such that: (A.9) ... (note: [34, Appendix H] also imposes some technical regularity conditions on the bounded set D for x in Eq. (1.1). For simplicity, we assume those conditions hold here.)"
Theorem 2.1 is the paper's central first-principles claim. Its proof does not establish the existence of the approximating SNN; instead it imports that existence from [34, Appendix H], a preprint by the same authors (Xia, Shen, Maini, Gaffney, Mogilner) that is not stated, proved, or machine-checked here. The new mixed-variable content only smooths the categorical coordinates in (A.2) and verifies inherited regularity properties; the decisive step (A.9) is exactly the continuous-variable universal approximation theorem of the authors' own prior work. The appended note then assumes away unstated technical conditions on D. Thus the central approximation theorem reduces to a self-citation chain: if [34, Appendix H] fails or its hypotheses are not met, Theorem 2.1 has no proof in this paper.
full rationale
I examined the claimed derivation chain for self-definitional reductions, fitted inputs presented as predictions, and renamed known results. No empirical benchmark reduces to a fitted constant: the loss function in Eq. (3.1) is an empirical surrogate for the proposed generalized W2 distance, and the tests use independent permutation tests, R2 statistics, and classification accuracy rather than the training loss itself. The generalization bound in Theorem 2.2 is derived from external quantitative bounds (Fournier-Guillin, the triangle inequality for Wasserstein distances) and does not assume its own conclusion. The one load-bearing circularity is Theorem 2.1. The proof's only mechanism for producing an SNN is the citation to [34, Appendix H], the authors' own prior preprint, which is never reproduced. Moreover, the proof explicitly assumes unstated 'technical regularity conditions' on the bounded set D, converting the stated theorem into a conditional statement whose hypotheses are not fully checked. Because the central theoretical result is forced through this self-citation chain, the circularity score is 8 rather than a lower score. The independent numerical content in Section 3 is not circular, but it does not rescue the theorem's derivation from its reliance on the unverified external self-citation.
Assumptions & free parameters
free parameters (3)
- λ (weight for continuous components in Eq. 2.1) =
Σ Var[y_i] for Example 3.3; time-averaged variance for Example 3.4; not used in 3.1-3.2
- δ (neighborhood radius in loss Eq. 2.20/3.1) =
0.025, 0.5√8, 0.3√7, 0.02 per example
- coefficient 4 in Ŵδ (Eq. 2.2) =
4
assumptions (4)
- domain assumption Assumption 2.1: uniform boundedness of y_x and Ŵ_y_x; independence of ω and W̆ω from x; uniform Lipschitz continuity of f_x and Ŵ_f_x under generalized W2; finiteness of mixed Sobolev norms; uniform continuity of f_x in continuous components.
- domain assumption Universal approximation theorem for SNNs on continuous random fields ([34, Appendix H])
- standard math Fournier-Guillin convergence rate for empirical measures in Wasserstein distance ([6, Theorem 1])
- standard math Triangle inequality for (generalized) Wasserstein distances ([4, Proposition 2.1])
Cite this review
Pith. "Pith review of A generalized Wasserstein-2 distance approach for efficient reconstruction of random field models using stochastic neural networks." pith.science (2026). https://pith.science/paper/KK5MY6CB
@misc{pith2026250705143,
author = {Pith},
title = {Pith review of: A generalized Wasserstein-2 distance approach for efficient reconstruction of random field models using stochastic neural networks},
year = {2026},
howpublished = {\url{https://pith.science/paper/KK5MY6CB}},
note = {Machine review of arXiv:2507.05143}
}
read the original abstract
In this work, we propose a novel generalized Wasserstein-2 distance approach for efficiently training stochastic neural networks to reconstruct random field models, where the target random variable comprises both continuous and categorical components. We prove that a stochastic neural network can approximate random field models under a Wasserstein-2 distance metric under nonrestrictive conditions. Furthermore, this stochastic neural network can be efficiently trained by minimizing our proposed generalized local squared Wasserstein-2 loss function. We showcase the effectiveness of our proposed approach in various uncertainty quantification tasks, including classification, reconstructing the distribution of mixed random variables, and learning complex noisy dynamical systems from spatiotemporal data.
Figures
Figures from the paper (2 more)
Reference graph
Works this paper leans on
-
[34]
M. Xia, Q. Shen, P. Maini, E. Gaffney, and A. Mogilner , A new local time-decoupled squared wasserstein-2 method for training stochastic neural netwo rks to reconstruct uncertain parameters in dynamical systems, arXiv preprint arXiv:2503.05068, (2025)
arXiv 2025
-
[1]
Agresti , An Introduction to Categorical Data Analysis , Wiley, 3rd ed., 2018
A. Agresti , An Introduction to Categorical Data Analysis , Wiley, 3rd ed., 2018
work page 2018
-
[2]
G. Arya, M. Schauer, F. Sch ¨afer, and C. Rackauckas , Automatic differentiation of programs with discrete randomness, Adv. Neural Inf. Process., 35 (2022), pp. 10435–10447
work page 2022
-
[3]
D. J. Bartholomew, M. Knott, and I. Moustaki , Latent Variable Models and Factor Analysis , Wiley, 7th ed., 2011
work page 2011
-
[4]
P. Clement and W. Desch , An elementary proof of the triangle inequality for the Wasse rstein metric , Proc. Amer. Math. Soc., 136 (2008), pp. 333–339
work page 2008
-
[5]
Cuturi , Sinkhorn distances: lightspeed computation of optimal tra nsport, Adv
M. Cuturi , Sinkhorn distances: lightspeed computation of optimal tra nsport, Adv. Neural Inf. Process., 26 (2013)
work page 2013
-
[6]
N. Fournier and A. Guillin , On the rate of convergence in Wasserstein distance of the emp irical measure, Probab. Theory Relat. Fields, 162 (2015), pp. 707–738
work page 2015
-
[7]
L. M. Friedman, C. D. Furberg, D. L. DeMets, D. M. Reboussin, a nd C. B. Granger , Funda- mentals of clinical trials , Springer, 2015
work page 2015
Show all 35 references
-
[8]
Gao and M
Y. Gao and M. K. Ng , Wasserstein generative adversarial uncertainty quantific ation in physics-informed neural networks, J. Comput. Phys., 463 (2022), p. 111270
2022
-
[9]
Ghahramani , Bayesian nonparametrics and the probabilistic approach to modelling, Phil
Z. Ghahramani , Bayesian nonparametrics and the probabilistic approach to modelling, Phil. Trans. R. Soc. A, 371 (2013), p. 20110553. This manuscript is for review purposes only. A GENERALIZED WASSERSTEIN-2 DISTANCE APPROACH 23
2013
-
[10]
Ghahramani and M
Z. Ghahramani and M. Jordan , Supervised learning from incomplete data via an EM approach , Adv. Neural Inf. Process., 6 (1993)
1993
-
[11]
Hastie, R
T. Hastie, R. Tibshirani, and J. Friedman , The Elements of Statistical Learning , Springer, 2nd ed., 2009
2009
-
[12]
K. He, X. Zhang, S. Ren, and J. Sun , Deep residual learning for image recognition , in CVPR, 2016, pp. 770–778
2016
-
[13]
Jaruszewicz and T
J. Jaruszewicz and T. Lipniacki , Toggle switch: noise determines the winning gene , Phys. Biol., 10 (2013), p. 035007
2013
-
[14]
Q. Jin, X. Luo, Y. Shi, and K. Kita , Image generation method based on improved condition GAN , in ICSAI, IEEE, 2019, pp. 1290–1294
2019
-
[15]
Koller and N
D. Koller and N. Friedman , Probabilistic Graphical Models: Principles and Technique s, MIT Press, 2009
2009
-
[16]
Lakshminarayanan, A
B. Lakshminarayanan, A. Pritzel, and C. Blundell , Simple and scalable predictive uncertainty estimation using deep ensembles , in Adv. Neural Inf. Process., vol. 30, 2017
2017
-
[17]
C. E. McCulloch, S. R. Searle, and J. M. Neuhaus , Generalized, Linear, and Mixed Models , Wiley, 2nd ed., 2008
2008
-
[18]
MCF ADDEN, Conditional logit analysis of qualitative choice behavior , Front
D. MCF ADDEN, Conditional logit analysis of qualitative choice behavior , Front. Econom., (1974)
1974
-
[19]
Mullachery, A
V. Mullachery, A. Khera, and A. Husain , Bayesian neural networks , arXiv preprint arXiv:1801.07710, (2018)
2018 arXiv
-
[20]
W. Nash, T. Sellers, S. Talbot, A. Cawthorn, and W. Ford , Abalone. UCI Machine Learning Repository, 1994. DOI: https://doi.org/10.24432/C55C7W
1994 doi
-
[21]
A. S. Now ak and K. R. Collins , Reliability of Structures , CRC press, 2012
2012
-
[22]
NYU, NYU HPC hardware specs , 2025
2025
-
[23]
Pesarin and L
F. Pesarin and L. Salmaso , Permutation Tests for Complex Data: Theory, Applications a nd Software, John Wiley & Sons, 2010
2010
-
[24]
Schiebinger, J
G. Schiebinger, J. Shu, M. Tabaka, et al. , Optimal-transport analysis of single-cell gene expressio n identifies developmental trajectories in reprogramming , Cell, 176 (2019), pp. 928–943.e22
2019
-
[25]
Sensoy, L
M. Sensoy, L. Kaplan, and M. Kandemir , Evidential deep learning to quantify classification uncer- tainty, Adv. Neural Inf. Process., 31 (2018)
2018
-
[26]
B. W. Sil verman, Density Estimation for Statistics and Data Analysis , Chapman & Hall, 1986
1986
-
[27]
Solomon, F
J. Solomon, F. De Goes, G. Peyr ´e, M. Cuturi, A. Butscher, A. Nguyen, T. Du, and L. Guibas , Convolutional Wasserstein distances: Efficient optimal tra nsportation on geometric domains , ACM Trans. Graph., 34 (2015), pp. 1–11
2015
-
[28]
Tian and K
T. Tian and K. Burrage , Stochastic models for regulatory networks of the genetic to ggle switch , Proc. Natl. Acad. Sci., 103 (2006), pp. 8372–8377
2006
-
[29]
Villani et al
C. Villani et al. , Optimal Transport: Old and New , vol. 338, Springer, Heidelberg, 2009
2009
-
[30]
W ang, J
J. W ang, J. Wu, X. Huang, and Z. Xiong , Improved WGAN for image generation methods , in ICMNM, Springer, 2023, pp. 199–211
2023
-
[31]
M. Xia, X. Li, Q. Shen, and T. Chou , An efficient Wasserstein-distance approach for reconstruct ing jump-diffusion processes using parameterized neural networ ks, Mach. Learn.: Sci. Technol., 5 (2024), p. 045052
2024
-
[32]
, Squared Wasserstein-2 distance for efficient reconstructio n of stochastic differential equations , arXiv preprint arXiv:2401.11354, (2024)
2024 arXiv
-
[33]
Xia and Q
M. Xia and Q. Shen , A local squared Wasserstein-2 method for efficient reconstru ction of models with uncertainty, arXiv preprint arXiv:2406.06825, (2024)
2024 arXiv
-
[35]
Zheng, F.-Y
W. Zheng, F.-Y. W ang, and C. Gou , Nonparametric different-feature selection using Wasserste in distance, in ICTAI, IEEE, 2020, pp. 982–988. This manuscript is for review purposes only
2020
Reviewed August 6, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.