Deep Kuratowski Embedding Neural Networks for Wasserstein Metric Learning
This provides a faster surrogate for Wasserstein distance computations, which is incremental for data analysis pipelines relying on such metrics.
The paper tackled the computational bottleneck of pairwise Wasserstein distances by proposing neural architectures to approximate the Wasserstein-2 distance, with ODE-KENN achieving a 28% lower test MSE than a baseline and 18% lower than another proposed method on MNIST data.
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.