Deep Kuratowski Embedding Neural Networks for Wasserstein Metric Learning

Andrew Qing He

cs.LGApr 6, 2026

Deep Kuratowski Embedding Neural Networks for Wasserstein Metric Learning

Authors: Andrew Qing He

Organizations: Department of Mathematics, Southern Methodist University

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.

Figures & tables

Model	Test MSE	Test MAE	Rel. MAE
Naive (single-layer)	$4.67\times 10^{-3}$	$5.249\times 10^{-2}$	1.63%
DeepKENN	$4.07\times 10^{-3}$	$4.873\times 10^{-2}$	1.52%
ODE-KENN	$\mathbf{3.35\times 10^{-3}}$	$\mathbf{4.406\times 10^{-2}}$	1.37%

Table 1: Test set performance on MNIST

W_{2}

distance learning (2,750 pairs). Distances are in pixel units. Rel. MAE = MAE / mean

(W_{2})

Figure 1: Experimental results for all three models trained for 2,000 epochs on MNIST

W_{2}

distance learning. Top left: Training (dashed) and validation (solid) MSE loss curves on a log scale. ODE-KENN (green) converges fastest and maintains the smallest train/val gap throughout training. Top right: Predicted vs. true

W_{2}

distances on the test set (2,750 pairs). ODE-KENN produces the tightest scatter around the ideal diagonal, particularly at small

W_{2}

values. Bottom left: Learned layer weights

\lambda_{k}

for DeepKENN. The first two convolutional layers are nearly fully suppressed (

\lambda\approx 0

), while FC1 dominates (

\lambda\approx 1.62

), indicating that abstract compressed representations are most informative for

W_{2}

. Bottom right: Learned time weights

\lambda_{t}

for ODE-KENN over

t\in[0,1]

. The bell-shaped profile peaks near

t\approx 0.35

, downweighting the raw feature (

t=0

) and the late trajectory, and assigning greatest importance to the early-to-mid ODE evolution.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Layer	Operation	Output shape	Flat dim
Conv1	Conv $(1\to 8,5\times 5)$ + ReLU + MaxPool $(2)$	$(8,14,14)$	1568
Conv2	Conv $(8\to 16,3\times 3)$ + ReLU + MaxPool $(2)$	$(16,7,7)$	784
Conv3	Conv $(16\to 32,3\times 3)$ + ReLU + MaxPool $(2)$	$(32,3,3)$	288
FC1	Linear $(288\to 128)$ + ReLU	$(128,)$	128
FC2	Linear $(128\to 64)$	$(64,)$	64

Appendix

Table 2: CNN encoder architecture shared by all three models.

Wasserstein Distance

Wasserstein distance compares probability distributions by the minimum cost of transporting mass between them.

See all 61 papers

Deep Residual Networks Learn the Geodesic Curve in the Wasserstein Space
Feb 18, 2021Kuo Gai, Shihua ZhangResidual Learning Neural Network Training Dynamics
Shanghai Institute for Mathematics and Interdisciplinary Sciences (SIMIS) Shanghai 200433, China · Academy of Mathematics and Systems Science Chinese Academy of Sciences Beijing 100190, China · School of Mathematical Sciences2026 University of Chinese Academy of Sciences Beijing 100049, China
Linearized Wasserstein dimensionality reduction with approximation guarantees
Feb 14, 2023Alexander Cloninger, Keaton Hamm, Varun Khurana +1Dimensionality Reduction Manifold Learning
DEPARTMENT OF MATHEMATICS, UNIVERSITY OF CALIFORNIA, SAN DIEGO, CA · HALICIOĦGLU DATA SCIENCE INSTITUTE, UNIVERSITY OF CALIFORNIA, SAN DIEGO, CA · DEPARTMENT OF MATHEMATICS, UNIVERSITY OF TEXAS AT ARLINGTON, ARLINGTON, TX +1
$k$ -Nearest Neighbors in Gromov--Wasserstein Space
Jun 9, 2026Kaitlyn Hohmeier, Nicolas Fraiman, Caroline MoosmuellerGraph Classification k-Nearest Neighbors
University of North Carolina at Chapel Hill, Department of Mathematics · University of North Carolina at Chapel Hill, Department of Statistics and Operations Research

Feb 18, 2021J/K move · Enter open · S save

Deep Residual Networks Learn the Geodesic Curve in the Wasserstein Space

Kuo Gai, Shihua Zhang

Shanghai Institute for Mathematics and Interdisciplinary Sciences (SIMIS) Shanghai 200433, China · Academy of Mathematics and Systems Science Chinese Academy of Sciences Beijing 100190, China · School of Mathematical Sciences2026 University of Chinese Academy of Sciences Beijing 100049, China

Residual Learning Neural Network Training Dynamics Optimal Transport

Open paper AskPDF arXiv

Recent studies revealed the mathematical connection between deep neural networks (DNNs) and dynamic systems. However, the specific dynamics that DNNs, especially deep residual networks (ResNets), tend to learn during training remain insufficiently characterized. To this end, we model the forward propagation of deep residual networks using continuity equations, in which the measure is conserved and infinite curves in the measure space connect the input distribution to the output one of a ResNet. We find ResNets with

L_2

regularization attempt to learn the geodesic curve in the Wasserstein space, induced by the optimal transport map. Compared with plain networks, ResNets can better approximate the geodesic curve, which explains why ResNets can be optimized and generalize better. Numerical experiments show that the data tracks of a ResNet tend to be line-shaped in terms of the line-shape score, and the map learned by a ResNet is closer to the optimal transport map in terms of the optimal transport score. In a word, we conclude that ResNets learn the geodesic curve in the Wasserstein space and discretely engineer the data transformation in high-dimensional spaces.

Abstract

Wasserstein Distance

Explore similar work

Deep Residual Networks Learn the Geodesic Curve in the Wasserstein Space

Linearized Wasserstein dimensionality reduction with approximation guarantees

kkk-Nearest Neighbors in Gromov--Wasserstein Space

$k$ -Nearest Neighbors in Gromov--Wasserstein Space