OTAKNet: Optimal Transport Adaptive Kalman Filter

• OTAKNet is a learning-based adaptive Kalman filter that combines an offline-trained KalmanNet with an optimal transport module for online, label-free noise drift adaptation.
• It aligns the one-step predictive measurement likelihood with an innovation-informed target distribution using Wasserstein-2 distance for stable, differentiable updates.
• Empirical evaluations show OTAKNet outperforms traditional adaptive filters and adapts within ~25 steps, both in synthetic Lorenz systems and real-world robot localization.

OTAKNet, short for Optimal-Transport-based Adaptive Kalman Network, is a learning-based adaptive Kalman filtering method designed for online, label-free adaptation to noise-statistics drift in nonlinear dynamical systems (He et al., 9 Aug 2025). It is introduced in “Differentiable Adaptive Kalman Filtering via Optimal Transport” as a method that wraps an offline-trained KalmanNet backbone with a differentiable optimal-transport adaptation loop operating at test time (He et al., 9 Aug 2025). The central premise is that drift in process and measurement noise covariances alters the one-step predictive measurement likelihood, and that this mismatch can be corrected online by aligning predictive and observation-derived measurement distributions in Wasserstein space, without ground-truth states and without retraining on labeled target-domain trajectories (He et al., 9 Aug 2025).

1. Problem setting and conceptual scope

OTAKNet is formulated for nonlinear state-space models with Gaussian process and measurement noise:

$$x_t = \mathbf{f}(x_{t-1}) + w_t,\qquad w_t \sim \mathcal{N}(0,\mathbf{Q}),$$

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$

where $x_t \in \mathbb{R}^n$ is the latent state, $y_t \in \mathbb{R}^m$ is the observation, $\mathbf f$ is the state-transition function, $\mathbf h$ is the observation function, and $\mathbf Q,\mathbf R$ are process and measurement noise covariance matrices (He et al., 9 Aug 2025). The method addresses noise-statistics drift, meaning a mismatch between the test-time true noise covariances and the covariances assumed by the filter (He et al., 9 Aug 2025).

The paper places OTAKNet at the intersection of adaptive Kalman filtering, learning-based filtering, and test-time adaptation. Its motivation is that classical environmental and operational factors—such as changing wind conditions or electromagnetic interference—can alter effective $\mathbf Q$ and $\mathbf R$, and that learning-based filters trained offline under fixed or sampled noise regimes can degrade substantially under such drift (He et al., 9 Aug 2025). The paper further argues that prior learning-based approaches mainly compensate by enlarging training data, training across many noise ratios, or performing offline fine-tuning on labeled target trajectories, none of which constitutes true online unlabeled adaptation (He et al., 9 Aug 2025).

The derivation uses EKF-like local Gaussian approximations through Jacobians $\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$ and $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$0, but OTAKNet is not introduced as a pure EKF or UKF variant. The learned backbone is explicitly an offline-trained KalmanNet, and OTAKNet adds an OT-based online adaptation mechanism around that learned filter (He et al., 9 Aug 2025).

2. Predictive-likelihood alignment via optimal transport

The core idea is to adapt the state estimator through the one-step predictive measurement likelihood

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$1

Under the local Gaussian approximation used in the paper,

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$2

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$3

which yields

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$4

with

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$5

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$6

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$7

The paper’s claim is that drift changes this predictive measurement distribution, so aligning it with an observation-derived target distribution provides a label-free adaptation signal (He et al., 9 Aug 2025).

The adaptation objective is written as

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$8

where $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$9 is the Wasserstein-2 distance in measurement space (He et al., 9 Aug 2025). The paper contrasts this with pointwise measurement matching of the form

$x_t \in \mathbb{R}^n$0

and argues that pointwise online adaptation is inadequate because a single observation yields high-variance gradients, single-step matching lacks memory of drift, residuals conflate model bias and noise effects, and independent per-step matching breaks the recursive temporal structure of filtering (He et al., 9 Aug 2025).

Optimal transport is used because, in the paper’s formulation, it is geometry-aware, can compare empirical distributions even when supports are shifted, provides stable, non-degenerate gradients, and remains differentiable under regularization (He et al., 9 Aug 2025). The method therefore recasts adaptive filtering under drift as a distributional alignment problem in predictive measurement space rather than a direct covariance-matching problem alone.

3. Source and target distributions, OT loss, and online update

A defining design choice in OTAKNet is the construction of empirical source and target distributions. The source distribution approximates the predictive measurement likelihood produced by the current adapted model. Given the current prior mean and predictive measurement covariance, the method samples

$x_t \in \mathbb{R}^n$1

and forms the empirical measure

$x_t \in \mathbb{R}^n$2

This is the measurement-space predictive distribution under the current filter parameters (He et al., 9 Aug 2025).

The target distribution is built from a sliding window of recent innovations,

$x_t \in \mathbb{R}^n$3

followed by the pseudo-measurements

$x_t \in \mathbb{R}^n$4

and the empirical target measure

$x_t \in \mathbb{R}^n$5

The target therefore encodes recent innovation history around the current observation rather than relying on the single measurement $x_t \in \mathbb{R}^n$6 alone (He et al., 9 Aug 2025). The paper presents this as the mechanism by which OTAKNet injects short-horizon temporal memory into online drift adaptation.

The OT problem is constructed with cost matrix

$x_t \in \mathbb{R}^n$7

uniform marginals

$x_t \in \mathbb{R}^n$8

and regularized coupling

$x_t \in \mathbb{R}^n$9

The paper uses IPOT rather than vanilla Sinkhorn. It defines

$y_t \in \mathbb{R}^m$0

then iteratively updates

$y_t \in \mathbb{R}^m$1

$y_t \in \mathbb{R}^m$2

and

$y_t \in \mathbb{R}^m$3

The adaptation loss is

$y_t \in \mathbb{R}^m$4

and the network parameters are updated online by

$y_t \in \mathbb{R}^m$5

All operations are implemented in PyTorch, allowing differentiation of the OT loss with respect to the neural filter parameters $y_t \in \mathbb{R}^m$6 (He et al., 9 Aug 2025).

The paper also introduces a warm-up learning-rate schedule to stabilize early adaptation when the residual buffer is still short:

$y_t \in \mathbb{R}^m$7

This schedule is explicitly motivated by the poor estimation of the target distribution during the first few steps (He et al., 9 Aug 2025).

4. Architectural composition and algorithmic workflow

Architecturally, OTAKNet consists of six components: an offline-trained KalmanNet backbone, a predictive-likelihood sampling module, a residual-memory/sliding-window module, a target-distribution constructor, an OT solver layer, and a differentiable online updater (He et al., 9 Aug 2025). The adapted object is the neural filter parameter vector $y_t \in \mathbb{R}^m$8, not merely a pair of analytically updated covariance matrices. From the pseudocode and text, $y_t \in \mathbb{R}^m$9 outputs at least $\mathbf f$0, $\mathbf f$1, $\mathbf f$2, and $\mathbf f$3, or corresponding estimated quantities (He et al., 9 Aug 2025).

The online procedure is per time step. Initialization loads the offline-trained network parameters $\mathbf f$4, initializes a residual queue of capacity $\mathbf f$5, and produces

$\mathbf f$6

For each subsequent step $\mathbf f$7, the method performs $\mathbf f$8 OT-based gradient steps. In each such step it computes the prior prediction

$\mathbf f$9

$\mathbf h$0

and the predictive measurement covariance

$\mathbf h$1

It then samples source particles, computes the current innovation

$\mathbf h$2

stores $\mathbf h$3 in the residual queue, constructs the target support points

$\mathbf h$4

forms the empirical measures, solves the OT problem, computes $\mathbf h$5, updates $\mathbf h$6, and finally runs the adapted network forward:

$\mathbf h$7

The paper therefore describes OTAKNet not as a new closed-form Kalman recursion, but as a learned Kalman filter plus an online OT-based test-time adaptation loop (He et al., 9 Aug 2025).

This structural distinction is important for positioning the method relative to prior work. Against classical model-based adaptive KF, OTAKNet adapts a neural filter parameterization rather than directly updating $\mathbf h$8 and $\mathbf h$9 via covariance matching, variational Bayesian inference, innovation statistics, or smoothing rules. Against offline learning-based adaptive filtering, it performs true online adaptation during deployment with no ground-truth state labels, no offline target-domain retraining, and no need to know noise-ratio regimes in advance (He et al., 9 Aug 2025).

5. Empirical evaluation on Lorenz and NCLT

The synthetic benchmark is a Lorenz attractor model with

$\mathbf Q,\mathbf R$0

$\mathbf Q,\mathbf R$1

identity $\mathbf Q,\mathbf R$2, Gaussian noises, and $\mathbf Q,\mathbf R$3 (He et al., 9 Aug 2025). The noise setting is

$\mathbf Q,\mathbf R$4

Performance is reported as MSE in dB across drift levels $\mathbf Q,\mathbf R$5 dB (He et al., 9 Aug 2025). The baselines are Oracle EKF / OKF with true covariance values, VBAKF-PR, SWVAKF, AKNet, MAML-KalmanNet, KalmanNet, and OTAKNet (He et al., 9 Aug 2025).

The paper reports that OTAKNet consistently outperforms all non-oracle baselines across all tested Lorenz drift levels. Representative results include: at $\mathbf Q,\mathbf R$6 dB drift, KalmanNet $\mathbf Q,\mathbf R$7, OTAKNet $\mathbf Q,\mathbf R$8, SWVAKF $\mathbf Q,\mathbf R$9, AKNet $\mathbf Q$0, and MAML-KalmanNet $\mathbf Q$1; at $\mathbf Q$2 dB drift, KalmanNet $\mathbf Q$3, OTAKNet $\mathbf Q$4, SWVAKF $\mathbf Q$5, AKNet $\mathbf Q$6, and MAML-KalmanNet $\mathbf Q$7; and at $\mathbf Q$8 dB drift, KalmanNet $\mathbf Q$9 versus OTAKNet $\mathbf R$0 (He et al., 9 Aug 2025). The oracle EKF remains best because it uses the true noise covariances, but OTAKNet narrows that gap more effectively than the adaptive baselines considered (He et al., 9 Aug 2025). Under $\mathbf R$1 dB drift, OTAKNet is reported to adapt within about 25 time steps and approach the performance of EKF with true covariances (He et al., 9 Aug 2025).

The real-world benchmark uses NCLT robot localization from IMU, odometry, and GPS. The state is

$\mathbf R$2

with dynamics

$\mathbf R$3

and measurement

$\mathbf R$4

The noise assumptions are

$\mathbf R$5

$\mathbf R$6

The trajectory is divided into 20 chunks of 200 steps: 13 train, 3 validation, and 4 test (He et al., 9 Aug 2025).

Two training regimes are reported: Full training with 13 labeled trajectories offline and Limited training with only 3 trajectories offline (He et al., 9 Aug 2025). The quantitative results are as follows.

Setting	Method	Result
Full training	EKF	$\mathbf R$7
Full training	VBAKF	$\mathbf R$8
Full training	SWVAKF	$\mathbf R$9
Full training	KalmanNet	$\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$0
Full training	OTAKNet	$\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$1
Limited training	KalmanNet	$\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$2
Limited training	OTAKNet	$\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$3

These results support the paper’s claim that OTAKNet improves over KalmanNet under both training regimes and is especially useful when offline training data are limited (He et al., 9 Aug 2025). The paper also notes that VBAKF outperforms fixed-noise EKF on NCLT, which suggests that real environmental drift is present in that dataset (He et al., 9 Aug 2025). Time-wise MSE curves show OTAKNet outperforming online baselines for most time steps, and the trajectory plots suggest better adaptation during abrupt maneuvers at the beginning and end of the route (He et al., 9 Aug 2025).

6. Ablations, runtime, theoretical justification, and limitations

The ablation study removes two components: the warm-up schedule and OT itself, the latter by letting the target distribution degenerate to a single point without windowed history (He et al., 9 Aug 2025). On the synthetic benchmark the reported results are OTAKNet $\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$4, without warm-up $\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$5, and without OT $\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$6. On NCLT the reported results are OTAKNet $\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$7, without warm-up $\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$8, and without OT $\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})$9 (He et al., 9 Aug 2025). These ablations support the claim that both the OT-based distributional loss and the warm-up schedule contribute to performance.

The paper’s principal formal justification appears in an appendix proposition on online innovation covariance adaptive consistency. If

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$00

where

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$01

and

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$02

then for Gaussians

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$03

Under innovation stationarity and ergodicity, differentiability of the OT loss, and a Polyak–Łojasiewicz condition, gradient descent converges to a $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$04 such that

$$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$05

The paper does not claim a general theorem equating likelihood and OT; rather, it argues that OT alignment of predictive and innovation-informed distributions yields consistency of the internal predictive covariance with the true innovation covariance (He et al., 9 Aug 2025).

The computational cost is nontrivial. The paper states that Sinkhorn-type OT can be solved in $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$06 per iteration and uses IPOT for practical differentiable computation (He et al., 9 Aug 2025). Reported inference times in the Lorenz setting are EKF $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$07 s, SWVAKF $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$08 s, VBAKF $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$09 s, and OTAKNet $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$10 s (He et al., 9 Aug 2025). The method is therefore much slower than EKF and SWVAKF and roughly comparable to VBAKF. A plausible implication is that OTAKNet is best matched to deployments where online unlabeled adaptation is important and moderate extra computation is acceptable.

The paper also makes several assumptions and caveats explicit. OTAKNet is tied to a Gaussian/local Gaussian formulation, assumes access to meaningful $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$11 and $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$12, relies on an offline-trained KalmanNet-style backbone, and uses a target-distribution construction $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$13 that is empirically motivated rather than uniquely derived (He et al., 9 Aug 2025). It is not aimed at general non-Gaussian or multimodal filtering, despite the broader applicability of OT in other contexts (He et al., 9 Aug 2025).

7. Terminological clarification and common misconceptions

A common point of confusion concerns the similarity of the name OTAKNet to other named neural architectures. The paper “ADC-Net: An Open-Source Deep Learning Network for Automated Dispersion Compensation in Optical Coherence Tomography” concerns ADC-Net, an OCT-specific encoder–decoder for automated dispersion compensation, and does not mention OTAKNet anywhere (Ahmed et al., 2022). ADC-Net is a retinal SD-OCT image compensation model based on a modified U-Net / FCN style architecture, whereas OTAKNet is a KalmanNet-based online adaptive filter for noise-statistics drift in nonlinear state-space estimation (Ahmed et al., 2022).

Another misconception is that OTAKNet is simply a new analytical Kalman update rule. The paper does not support that interpretation. OTAKNet is instead a differentiable online adaptation mechanism layered on top of a learned Kalman filter, with adaptation driven by OT alignment in measurement space (He et al., 9 Aug 2025). Similarly, it is not presented as merely updating $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$14 and $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$15 through closed-form adaptive-KF formulas; the method updates the broader neural filter parameterization $$y_t = \mathbf{h}(x_t) + v_t,\qquad v_t \sim \mathcal{N}(0,\mathbf{R}),$$16 (He et al., 9 Aug 2025).

The paper’s own positioning is correspondingly precise. Relative to classical adaptive KFs such as VBAKF-PR and SWVAKF, OTAKNet is a learning-based backbone with distributional test-time adaptation. Relative to AKNet and MAML-KalmanNet, it is an online, label-free alternative to offline or label-dependent adaptation procedures (He et al., 9 Aug 2025). In that sense, OTAKNet is best understood as a method for online domain adaptation of predictive likelihoods inside learned Kalman filtering, rather than as a standalone replacement for Kalman filtering as a paradigm.