ODE-KENN: Continuous Metric Learning

• ODE-KENN is a continuous-depth distance-learning architecture that leverages Neural ODEs to learn embedding trajectories for approximating the Wasserstein-2 distance.
• It combines a CNN encoder with a Neural ODE, integrating continuous latent flows to capture geometric discrepancies beyond terminal-state differences.
• Empirical results on MNIST demonstrate that ODE-KENN outperforms discrete baselines with lower test errors and a reduced generalization gap.

ODE-KENN is a continuous-depth distance-learning architecture introduced as the Neural ODE variant of the Deep Kuratowski Embedding Neural Network framework. Its purpose is to approximate the Wasserstein-2 distance $W_2$ with a fast learned surrogate, replacing repeated optimal transport solves in pipelines that require many pairwise distance evaluations. The method combines a CNN encoder with a Neural ODE, embeds each input into a trajectory in $C^1([0,1],\mathbb{R}^d)$, and defines the learned distance by a weighted trajectory discrepancy rather than by a terminal feature difference alone. On MNIST with exact precomputed $W_2$ distances, it reports lower test error than both a single-layer baseline and the discrete DeepKENN variant, with a smaller generalization gap (He, 6 Apr 2026).

1. Conceptual basis and nomenclature

ODE-KENN is motivated by the Kuratowski–Wojdysławski embedding theorem, which states that every bounded metric space can be embedded isometrically into a Banach space $\ell^\infty(M)$ as a closed subset of a convex set. The model adopts a trainable analogue of this idea: instead of predicting a distance directly, it learns an embedding $x \mapsto \phi(x)$ and recovers distance from a norm of embedding differences. In the discrete formulation, DeepKENN aggregates discrepancies across intermediate CNN feature maps; in the continuous formulation, ODE-KENN replaces the layer stack with a Neural ODE and compares entire latent trajectories (He, 6 Apr 2026).

Within that formulation, ODE-KENN is specifically a metric-learning architecture for $W_2$ approximation rather than a generic label for all ODE-based neural-symbolic systems. This distinction matters because the acronym “KENN” also appears in a different literature as “Knowledge Enhanced Neural Networks,” a framework that injects logical clauses through residual knowledge-enhancement layers. That relational KENN work explicitly states that it does not discuss continuous-depth formulations, neural ODEs, or any ODE view of KENN (Daniele et al., 2022). Thus, despite the shared suffix, ODE-KENN and KENN denote different constructions.

A broader, informal reading is nevertheless possible. Several papers describe continuous-depth models with structured or interpretable differential equations in terms that are close to an “ODE-KENN-style” philosophy. In holographic QCD, for example, the Neural ODE weight function is identified with a bulk metric and the learned differential equation is interpreted as reconstructed geometry (Hashimoto et al., 2020). This suggests that ODE-KENN can also be situated within a wider family of ODE-based structured representations, even though its immediate target is Wasserstein metric approximation.

2. Mathematical formulation

The target quantity is the Wasserstein-2 distance

$$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$

where $\Pi(\mu,\nu)$ is the set of couplings and $c(x,y)$ is the ground cost. For MNIST, each image is treated as a discrete probability measure over pixels, with $C_{ij}=\|p_i-p_j\|_2^2$ (He, 6 Apr 2026).

DeepKENN, the discrete precursor, uses a CNN encoder

$C^1([0,1],\mathbb{R}^d)$0

and aggregates layerwise discrepancies as

$C^1([0,1],\mathbb{R}^d)$1

with $C^1([0,1],\mathbb{R}^d)$2, ensuring $C^1([0,1],\mathbb{R}^d)$3. This gives a Euclidean norm in the product space $C^1([0,1],\mathbb{R}^d)$4 with weighted inner product $C^1([0,1],\mathbb{R}^d)$5 (He, 6 Apr 2026).

ODE-KENN replaces the finite layer sequence by a continuous latent flow. Given an encoder $C^1([0,1],\mathbb{R}^d)$6, it defines a trajectory $C^1([0,1],\mathbb{R}^d)$7 by

$C^1([0,1],\mathbb{R}^d)$8

where $C^1([0,1],\mathbb{R}^d)$9 is an autonomous MLP vector field. The distance surrogate is

$W_2$0

with $W_2$1 a learnable time-dependent weight function, discretized at solver time steps and parameterized with softplus (He, 6 Apr 2026).

The continuous embedding space is $W_2$2. The paper notes that if the vector field is Lipschitz, then distinct initial conditions yield distinct trajectories, so the map is injective under reasonable conditions. It also gives the bound

$W_2$3

for an $W_2$4-Lipschitz vector field, and interprets the resulting smoothness as an implicit regularizer acting on trajectory space rather than on a finite collection of layers (He, 6 Apr 2026).

3. Architecture and training procedure

The practical ODE-KENN implementation uses a shared CNN encoder, followed by an RK4-integrated Neural ODE solved with torchdiffeq using $W_2$5 fixed steps over $W_2$6. To reduce computation, the two inputs in a pair are stacked into one batch and evolved together in a single ODE solve instead of solving two separate initial-value problems. The vector field is

$W_2$7

with $W_2$8 and $W_2$9, totaling 4,160 parameters (He, 6 Apr 2026).

All models in the reported experiment share the same CNN backbone:

• Conv $\ell^\infty(M)$0, $\ell^\infty(M)$1, ReLU, MaxPool
• Conv $\ell^\infty(M)$2, $\ell^\infty(M)$3, ReLU, MaxPool
• Conv $\ell^\infty(M)$4, $\ell^\infty(M)$5, ReLU, MaxPool
• FC $\ell^\infty(M)$6, ReLU
• FC $\ell^\infty(M)$7

The intermediate flattened dimensions are $\ell^\infty(M)$8, $\ell^\infty(M)$9, $x \mapsto \phi(x)$0, $x \mapsto \phi(x)$1, and $x \mapsto \phi(x)$2, so DeepKENN’s total embedding dimension is $x \mapsto \phi(x)$3 (He, 6 Apr 2026).

Training minimizes mean squared error to exact $x \mapsto \phi(x)$4: $x \mapsto \phi(x)$5 The MNIST dataset is converted into discrete probability measures by normalizing each $x \mapsto \phi(x)$6 grayscale image to sum to one. Exact $x \mapsto \phi(x)$7 distances were precomputed for 55,000 image pairs using ot.emd2 from POT, with 1,000 pairs for each of the 55 digit-class combinations. The split is 49,500 train, 2,750 validation, and 2,750 test. Training uses Adam, learning rate $x \mapsto \phi(x)$8, batch size 256, 2,000 epochs, and best-validation checkpoint selection (He, 6 Apr 2026).

To match parameter counts across baselines, the Naive and DeepKENN models are augmented with an extra FC$x \mapsto \phi(x)$9+Tanh head. The resulting counts are 55,424 for Naive, 55,430 for DeepKENN, and 55,434 for ODE-KENN (He, 6 Apr 2026).

4. Empirical performance and learned representations

The reported held-out performance on 2,750 test pairs is as follows (He, 6 Apr 2026):

Model	Test MSE	Test MAE
Naive	$W_2$0	$W_2$1
DeepKENN	$W_2$2	$W_2$3
ODE-KENN	$W_2$4	$W_2$5

The same evaluation reports relative MAE values of $W_2$6 for Naive, $W_2$7 for DeepKENN, and $W_2$8 for ODE-KENN. The paper summarizes these results as a 28% reduction in test MSE relative to the Naive baseline and an 18% reduction relative to DeepKENN, while also reporting the smallest generalization gap for the ODE-based model (He, 6 Apr 2026).

The learned weighting profiles provide an internal interpretation of where the surrogate extracts geometric information. In DeepKENN, the learned $W_2$9 values nearly suppress early convolutional layers: Conv1 $$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$0, Conv2 $$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$1, Conv3 $$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$2, FC1 $$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$3, and FC2 $$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$4. In ODE-KENN, the learned $$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$5 is bell-shaped and peaks around $$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$6, so the initial feature state and late-time states are comparatively downweighted. This suggests that, in the reported MNIST setting, the most informative geometry for $$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$7 approximation emerges in early-to-middle latent evolution rather than in raw input features or terminal states (He, 6 Apr 2026).

The paper further states that ODE-KENN’s scatter plots are closest to the diagonal, especially for small $$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$8 values. Because small pairwise distances are often important in manifold learning pipelines, this behavior is consequential for the stated downstream use case of replacing a Wasserstein oracle (He, 6 Apr 2026).

5. Relation to other ODE-based structured models

ODE-KENN belongs to a larger research pattern in which neural models are organized around continuous-time latent evolution rather than fixed stacks of layers. An MRI reconstruction model, for example, treats feature evolution as

$$W_2(\mu,\nu)=\left( \inf_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal X} c(x,y)^2\, d\gamma(x,y) \right)^{1/2},$$9

and also evolves parameters via

$\Pi(\mu,\nu)$0

while using an augmented state space to increase expressiveness (Yazdanpanah et al., 2019). A Hamiltonian graph simulator instead learns a Hamiltonian with a graph network and passes its induced vector field through a differentiable Runge–Kutta integrator, improving predictive accuracy, energy accuracy, and zero-shot generalization to unseen time-step sizes and integrator orders (Sanchez-Gonzalez et al., 2019).

Other ODE-based models emphasize structured interpretability more directly. In holographic QCD, the continuous weight $\Pi(\mu,\nu)$1 is identified with the bulk metric combination

$\Pi(\mu,\nu)$2

so learning the Neural ODE becomes equivalent to learning the bulk geometry from lattice QCD chiral-condensate data (Hashimoto et al., 2020). In generative modeling with known mechanistic dynamics, GOKU-net embeds a fixed ODE

$\Pi(\mu,\nu)$3

inside a VAE and learns only the inference maps and observation model, enabling recovery of physically meaningful hidden states and parameters (Linial et al., 2020). In dynamic graphs, TI-ODE decomposes the graph ODE vector field into multiple learnable interaction basis functions with time-dependent coefficients, addressing heterogeneous and time-varying interaction patterns (Wang et al., 27 Apr 2026).

These examples do not instantiate ODE-KENN itself. A plausible implication is that ODE-KENN is best viewed as one member of a broader continuous-depth design space in which the differential equation is not merely an optimization device but part of the model’s structural prior. What distinguishes ODE-KENN within that space is the specific Kuratowski-style metric-learning objective and the use of weighted trajectory discrepancy in $\Pi(\mu,\nu)$4 for approximating $\Pi(\mu,\nu)$5.

6. Numerical considerations, scope, and common confusions

Like other Neural ODE architectures, ODE-KENN inherits numerical issues associated with continuous-depth training. A separate line of work on proximal implicit ODE solvers argues that stiff neural ODEs make explicit adaptive solvers expensive, and proposes implicit schemes based on proximal inner–outer iterations to improve numerical stability and computational efficiency. That work validates advantages on continuous-depth graph neural networks and continuous normalizing flows, suggesting a potentially relevant solver direction whenever an ODE-KENN-like architecture becomes stiff in practice (Baker et al., 2022).

The term should also be disambiguated carefully. ODE-KENN is not the same as the relational KENN framework, whose mechanism is a differentiable residual layer that revises neural pre-activations according to logical clauses and learnable clause weights; that literature explicitly states that it does not discuss neural ODEs or continuous-depth formulations (Daniele et al., 2022). Nor is ODE-KENN related to unrelated uses of “ODE” as an acronym, such as Open-Set Dynamic Evaluation in multimodal hallucination benchmarking (Tu et al., 2024) or the ODE/IM correspondence in integrable models (Ito et al., 2022).

In its own intended sense, ODE-KENN designates a continuous Kuratowski-style embedding network for Wasserstein metric learning. Its defining move is to replace a discrete multi-layer embedding by a Neural ODE trajectory, thereby treating the comparison object not as a single feature vector but as an entire path in function space. The reported MNIST results indicate that, under matched parameter budgets, this continuous formulation yields lower test error and a smaller generalization gap than both a terminal-feature baseline and a discrete multi-layer aggregation scheme (He, 6 Apr 2026).