DeepKENN: A Metric-Learning Approach to Wasserstein-2

• DeepKENN is a supervised metric-learning architecture that approximates the Wasserstein-2 distance by aggregating weighted Euclidean distances across all CNN layers.
• It replaces expensive optimal transport solvers with a fast neural surrogate by embedding intermediate representations and computing a composite Euclidean norm.
• Empirical results on MNIST show a reduction in test MSE compared to a terminal-layer baseline, though performance lags behind ODE-KENN with a smaller train/validation gap.

DeepKENN is a supervised metric-learning architecture for approximating the Wasserstein-2 distance $W_2$ between image pairs by embedding each input into the collection of all intermediate representations of a convolutional neural network and defining the predicted distance as a positively weighted Euclidean norm across layers. It is introduced as a finite-dimensional, trainable analogue of the Kuratowski embedding viewpoint on metric spaces, with the practical objective of replacing expensive optimal transport computations in pairwise distance pipelines by a fast neural surrogate (He, 6 Apr 2026).

1. Definition and motivation

The central problem addressed by DeepKENN is the computational cost of pairwise Wasserstein distance evaluation. In the setting considered, many downstream pipelines require a large pairwise $W_2$ distance matrix, and exact computation for discrete measures is expensive because $W_2$ is obtained by solving an optimal transport problem. The reported cost for exact linear-program computation on $n$-point discrete measures is $O(n^3 \log n)$, and even approximate Sinkhorn methods remain costly at scale. On MNIST-sized images, pairwise matrix construction is described as the dominant cost in the pipeline (He, 6 Apr 2026).

DeepKENN is motivated by the classical Kuratowski embedding theorem, recalled in the paper as the statement that every bounded metric space embeds isometrically into a Banach space, specifically $\ell^\infty(M)$. The architecture does not claim to realize an exact Kuratowski embedding or an isometric embedding into $\ell^\infty$. Instead, it uses that theorem as conceptual motivation for learning a finite-dimensional embedding in which Euclidean-type distances approximate the target metric $W_2$ (He, 6 Apr 2026).

This design choice places DeepKENN in a specific part of the metric-learning landscape. It is not presented as a Siamese network with a separate learned pairwise head. Rather, the predicted distance is induced directly by a family of per-layer embeddings, aggregated with learned nonnegative weights. This suggests that the architecture is intended to preserve metric structure through the representation itself rather than through a post hoc regression head.

2. Architectural formulation

The backbone is a feedforward CNN with $L$ stages, denoted $F_{\boldsymbol{\vartheta}}$, understood as a composition of layer maps. The intermediate representation at layer $W_2$0 is denoted $W_2$1, and the equivalent notation $W_2$2 is explicitly identified in the source material (He, 6 Apr 2026).

The single-layer baseline uses only the terminal feature:

$W_2$3

DeepKENN replaces that terminal-feature metric by an aggregation over all intermediate features:

$W_2$4

or equivalently, in layer notation,

$W_2$5

The shared encoder specified in the appendix contains five feature levels:

Stage	Definition	Flat dimension
Conv1	$W_2$6	1568
Conv2	$W_2$7	784
Conv3	$W_2$8	288
FC1	$W_2$9	128
FC2	$W_2$0	64

Accordingly, the feature dimensions are

$W_2$1

with total product-space dimension

$W_2$2

Intermediate CNN feature maps are taken after each stage, flattened if necessary, and compared by Euclidean distance. Each layer contributes a scalar-weighted squared distance term. The aggregation weights are constrained to be nonnegative,

$W_2$3

and are implemented via softplus reparameterization,

$W_2$4

for unconstrained $W_2$5 (He, 6 Apr 2026).

3. Metric interpretation and formal properties

The paper gives DeepKENN an explicit product-space interpretation. The distance corresponds to Euclidean distance in

$W_2$6

equipped with weighted inner product

$W_2$7

If one defines the combined embedding

$W_2$8

then

$W_2$9

This is the basis for calling DeepKENN the “discrete analog of the Kuratowski embedding” (He, 6 Apr 2026).

The strongest formal guarantee provided is structural. For any fixed backbone $n$0 and nonnegative $n$1, DeepKENN satisfies nonnegativity and the triangle inequality because it is Euclidean distance in a product feature space. The positivity constraint on the weights is therefore not merely numerical; it is what preserves norm structure in the aggregated embedding (He, 6 Apr 2026).

Positive definiteness is treated more cautiously. The paper states that

$n$2

depends on injectivity of the joint map

$n$3

which is said to “hold generically for sufficiently expressive networks,” but this is not proved. As a result, DeepKENN may fail to be a true metric if distinct inputs collapse under the learned representation (He, 6 Apr 2026).

The paper also does not prove that aggregating distances across layers improves approximation to $n$4. The theoretical intuition is instead geometric and capacity-based: the final layer alone gives only a $n$5-dimensional embedding, whereas concatenating all layer outputs yields the larger effective embedding dimension $n$6. This suggests a more expressive finite-dimensional surrogate, but the source does not establish approximation-error bounds in terms of network capacity or sample size.

4. Supervised learning setup and experimental protocol

The learning problem is fully supervised. Inputs are image pairs $n$7, with each image treated as a probability measure by normalizing pixel intensities to sum to one. Supervision comes from precomputed oracle $n$8 distances obtained exactly using the POT library’s ot.emd2. The model output is a scalar $n$9 trained to approximate the true Wasserstein distance $O(n^3 \log n)$0 (He, 6 Apr 2026).

Training uses mean squared error against the exact oracle:

$O(n^3 \log n)$1

where $O(n^3 \log n)$2 defines the nonnegative weights via $O(n^3 \log n)$3. No additional regularizer is specified for DeepKENN, and there is no mention of target normalization or output rescaling (He, 6 Apr 2026).

The experimental setup is centered on MNIST. Each $O(n^3 \log n)$4 grayscale image is normalized to a probability measure. The ground transport cost is defined by squared Euclidean distance between pixel coordinates,

$O(n^3 \log n)$5

with $O(n^3 \log n)$6. Exact $O(n^3 \log n)$7 distances are precomputed for 55,000 image pairs. Pair construction is balanced across digit class combinations, using 1,000 random pairs for each of the 55 class combinations $O(n^3 \log n)$8 with $O(n^3 \log n)$9. The resulting dataset is split into 49,500 training pairs, 2,750 validation pairs, and 2,750 test pairs (He, 6 Apr 2026).

Optimization uses Adam with learning rate $\ell^\infty(M)$0, batch size $\ell^\infty(M)$1, for $\ell^\infty(M)$2 epochs, and the best validation checkpoint is selected for final testing. For parameter-count matching, the Naive and DeepKENN models are augmented with an extra $\ell^\infty(M)$3 head containing 4,160 parameters, because ODE-KENN’s ODE vector field also has 4,160 parameters. The reported total parameter counts are 55,424 for Naive, 55,430 for DeepKENN, and 55,434 for ODE-KENN (He, 6 Apr 2026).

5. Empirical results and learned behavior

DeepKENN is evaluated against two comparators: the “Naive” terminal-feature Euclidean baseline and ODE-KENN. On the test set of 2,750 pairs, the reported results are as follows (He, 6 Apr 2026):

Model	Test MSE	Test MAE
Naive	$\ell^\infty(M)$4	$\ell^\infty(M)$5
DeepKENN	$\ell^\infty(M)$6	$\ell^\infty(M)$7
ODE-KENN	$\ell^\infty(M)$8	$\ell^\infty(M)$9

Relative MAE values are also reported: $\ell^\infty$0 for Naive, $\ell^\infty$1 for DeepKENN, and $\ell^\infty$2 for ODE-KENN. DeepKENN reduces test MSE from $\ell^\infty$3 to $\ell^\infty$4 relative to the single-layer baseline, which is described as about a 13% reduction. ODE-KENN achieves an additional 18% lower test MSE than DeepKENN under nearly identical parameter count, and 28% lower test MSE than the single-layer baseline (He, 6 Apr 2026).

The paper emphasizes that DeepKENN improves over the terminal-feature baseline but generalizes less well than ODE-KENN. From the training and validation loss curves, ODE-KENN has the smallest train/validation gap, whereas DeepKENN has a larger gap. The stated interpretation is that DeepKENN’s larger effective embedding dimension $\ell^\infty$5 increases expressiveness but also overfitting risk, while ODE-KENN benefits from implicit regularization via trajectory smoothness (He, 6 Apr 2026).

A distinctive empirical finding concerns the learned weight profile across layers. The approximate learned values are reported as:

• Conv1: $\ell^\infty$6
• Conv2: $\ell^\infty$7
• Conv3: $\ell^\infty$8
• FC1: $\ell^\infty$9
• FC2: $W_2$0

This indicates that early convolutional layers are heavily suppressed, while most of the useful metric information for MNIST $W_2$1 approximation is concentrated in the compressed fully connected representations. The paper interprets this as evidence that low-level spatial feature distances contribute little in this setting, whereas later semantic or compressed features carry most of the relevant structure (He, 6 Apr 2026). A plausible implication is that the value of DeepKENN lies less in indiscriminately aggregating all layers than in learning which representation levels are actually informative for the target metric.

6. Scope, limitations, and terminological boundaries

DeepKENN is presented as useful when exact Wasserstein distances are needed repeatedly after an upfront supervised training stage. Once trained, inference consists of a forward pass and Euclidean-distance computations across feature maps, allowing the surrogate to replace repeated optimal transport solves in downstream pairwise distance computations. This makes the method an amortized approximation strategy rather than a direct solver for optimal transport (He, 6 Apr 2026).

The limitations are explicit. DeepKENN requires a supervised dataset of exact $W_2$2 distances, which is itself costly to construct. It preserves the triangle inequality by construction but does not enforce injectivity of the learned joint feature map, so positive definiteness is not guaranteed. The paper also does not establish out-of-distribution robustness or provide a theorem bounding approximation error to the true Wasserstein metric (He, 6 Apr 2026).

Compared with ODE-KENN, DeepKENN is simpler in implementation because it does not require an ODE solver and remains a discrete-depth architecture directly tied to intermediate CNN features. However, under matched parameter counts it is empirically less accurate and exhibits a larger generalization gap. The paper therefore frames the choice pragmatically: prefer DeepKENN over exact OT when many future pairwise queries justify one-time supervised training, and prefer DeepKENN over ODE-KENN when implementation simplicity or avoiding ODE machinery matters more than achieving the strongest reported test error and generalization (He, 6 Apr 2026).

A terminological clarification is also important. DeepKENN should not be conflated with “DeepKAN” or “DeepKANSeg,” which arise in a different arXiv paper on remote sensing semantic segmentation and are based on Kolmogorov–Arnold Networks. That paper explicitly states that it does not define a module called DeepKENN; its relevant terms are DeepKAN, GLKAN, and DeepKANSeg (Ma et al., 13 Jan 2025). This distinction matters because the similarity in names is nominal rather than conceptual: DeepKENN is a Wasserstein metric-learning architecture motivated by Kuratowski embedding, whereas DeepKAN is a KAN-based feature-refinement component in an encoder–decoder segmentation model.