[2604.04343] Deep Kuratowski Embedding Neural Networks for Wasserstein Metric Learning

Computer Science > Machine Learning arXiv:2604.04343 (cs) [Submitted on 6 Apr 2026] Title:Deep Kuratowski Embedding Neural Networks for Wasserstein Metric Learning Authors:Andrew Qing He View a PDF of the paper titled Deep Kuratowski Embedding Neural Networks for Wasserstein Metric Learning, by Andrew Qing He View PDF HTML (experimental) Abstract:Computing pairwise Wasserstein distances is a fundamental bottleneck in data analysis pipelines. Motivated by the classical Kuratowski embedding theorem, we propose two neural architectures for learning to approximate the Wasserstein-2 distance ($W_2$) from data. The first, DeepKENN, aggregates distances across all intermediate feature maps of a CNN using learnable positive weights. The second, ODE-KENN, replaces the discrete layer stack with a Neural ODE, embedding each input into the infinite-dimensional Banach space $C^1([0,1], \mathbb{R}^d)$ and providing implicit regularization via trajectory smoothness. Experiments on MNIST with exact precomputed $W_2$ distances show that ODE-KENN achieves a 28% lower test MSE than the single-layer baseline and 18% lower than DeepKENN under matched parameter counts, while exhibiting a smaller generalization gap. The resulting fast surrogate can replace the expensive $W_2$ oracle in downstream pairwise distance computations. Subjects: Machine Learning (cs.LG) MSC classes: 68T07, 49Q22, 46B85 Cite as: arXiv:2604.04343 [cs.LG]   (or arXiv:2604.04343v1 [cs.LG] for this version)   https://doi.org...