---
title: 'OTAKNet: Optimal Transport Adaptive Kalman Filter'
url: https://www.emergentmind.com/topics/otaknet
type: topic
---

# OTAKNet: Optimal Transport Adaptive Kalman Filter

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 [2508.07037]. 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 [2508.07037]. 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 [2508.07037].

## 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 [2508.07037]. The method addresses **noise-statistics drift**, meaning a mismatch between the test-time true noise covariances and the covariances assumed by the filter [2508.07037].

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 [2508.07037]. 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 [2508.07037].

The derivation uses EKF-like local Gaussian approximations through Jacobians \(\mathbf{F}_t=\nabla_x \mathbf f(\hat x_{t-1})\) and \(\mathbf{H}_t=\nabla_x \mathbf h(\hat x_{t|t-1})\), 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 [2508.07037].

## 2. Predictive-likelihood alignment via optimal transport

The core idea is to adapt the state estimator through the **one-step predictive measurement likelihood**
$$
p(y_t\mid y_{1:t-1}) = \int p(y_t\mid x_t)\,p(x_t\mid y_{1:t-1})\,dx_t.
$$
Under the local Gaussian approximation used in the paper,
$$
p(x_t\mid y_{1:t-1}) = \mathcal{N}\bigl(x_t;\hat x_{t|t-1},\Sigma_{t|t-1}\bigr),
$$
$$
p(y_t\mid x_t) = \mathcal{N}\bigl(y_t;\mathbf{h}(x_t),\hat{\mathbf R}\bigr),
$$
which yields
$$
p(y_t\mid y_{1:t-1}) = \mathcal{N}\bigl(y_t;\mathbf{h}(\hat x_{t|t-1}),\mathbf S_{t|t-1}\bigr),
$$
with
$$
\hat x_{t|t-1}=\mathbf f(\hat x_{t-1}),
$$
$$
\Sigma_{t|t-1} = \mathbf F_t \Sigma_{t-1}\mathbf F_t^\top + \hat{\mathbf Q},
$$
$$
\mathbf S_{t|t-1} = \mathbf H_t \Sigma_{t|t-1}\mathbf H_t^\top + \hat{\mathbf R}.
$$
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 [2508.07037].

The adaptation objective is written as
$$
\theta_t^* = \arg\min_{\theta} W_2^2\!\bigl(p(y_t(\theta)\mid y_{1:t-1}),\;p(y_t\mid y_{1:t-1})\bigr),
$$
where \(W_2\) is the Wasserstein-2 distance in measurement space [2508.07037]. The paper contrasts this with pointwise measurement matching of the form
$$
\min_{\theta} \sum_{t=1}^T \mathbb{E}\bigl\|\hat y_t(\theta)-y_t\bigr\|^2,
$$
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 [2508.07037].

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 [2508.07037]. 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
$$
z_{\mathrm{src}^{(i)}} \sim \mathcal N\bigl(\mathbf h(\hat x_{t|t-1}(\theta)),\,S_{t|t-1}(\theta)\bigr), \qquad i=1,\dots,N,
$$
and forms the empirical measure
$$
\mu_{\mathrm{src}} = \frac{1}{N}\sum_{i=1}^N \delta_{z_{\mathrm{src}^{(i)}}}.
$$
This is the measurement-space predictive distribution under the current filter parameters [2508.07037].

The **target distribution** is built from a sliding window of recent innovations,
$$
e_k = y_k - \mathbf h(\hat x_{k|k-1}), \qquad k=t-W+1,\dots,t,
$$
followed by the pseudo-measurements
$$
\tilde z_{\mathrm{tgt}^{(j)}} = y_t + e_{t-W+j},\qquad j=1,\dots,W,
$$
and the empirical target measure
$$
\mu_{\mathrm{tgt}} = \frac{1}{W}\sum_{j=1}^W \delta_{\tilde z_{\mathrm{tgt}^{(j)}}}.
$$
The target therefore encodes recent innovation history around the current observation rather than relying on the single measurement \(y_t\) alone [2508.07037]. 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
$$
C_{ij} = \frac12 \left\| z_{\mathrm{src}^{(i)}}-\tilde z_{\mathrm{tgt}^{(j)}} \right\|^2,
$$
uniform marginals
$$
a_i=\frac1N,\qquad b_j=\frac1W,
$$
and regularized coupling
$$
\pi^* = \arg\min_{\pi\in\Pi(a,b)} \sum_{i=1}^N \sum_{j=1}^W \pi_{ij} C_{ij} -\varepsilon h(\pi).
$$
The paper uses **IPOT** rather than vanilla Sinkhorn. It defines
$$
G_{ij} = \exp(-C_{ij}/\epsilon),
$$
then iteratively updates
$$
Q = G \odot \pi^{(k)},
$$
$$
a \leftarrow \frac{\mu}{Qb}, \qquad b \leftarrow \frac{\nu}{Q^\top a},
$$
and
$$
\pi^{(k+1)} \leftarrow \operatorname{diag}(a)\,Q\,\operatorname{diag}(b).
$$
The adaptation loss is
$$
\mathcal L_{\mathrm{OT}} = \sum_{i=1}^N\sum_{j=1}^W \pi_{ij} C_{ij},
$$
and the network parameters are updated online by
$$
\theta \leftarrow \theta - \eta_t \nabla_\theta \mathcal L_{\mathrm{OT}}.
$$
All operations are implemented in PyTorch, allowing differentiation of the OT loss with respect to the neural filter parameters \(\theta\) [2508.07037].

The paper also introduces a warm-up learning-rate schedule to stabilize early adaptation when the residual buffer is still short:
$$
\eta_t = \min\!\left(\eta,\frac{t}{W}\eta\right).
$$
This schedule is explicitly motivated by the poor estimation of the target distribution during the first few steps [2508.07037].

## 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** [2508.07037]. The adapted object is the neural filter parameter vector \(\theta\), not merely a pair of analytically updated covariance matrices. From the pseudocode and text, \(OTAKNet_\theta\) outputs at least \(\hat x_t\), \(\Sigma_{t|t}\), \(\mathbf Q_t\), and \(\mathbf R_t\), or corresponding estimated quantities [2508.07037].

The online procedure is per time step. Initialization loads the offline-trained network parameters \(\theta\), initializes a residual queue of capacity \(W\), and produces
$$
\hat x_1(\theta),\ \Sigma_{1|1}(\theta),\ \hat{\mathbf Q}_1(\theta),\ \hat{\mathbf R}_1(\theta).
$$
For each subsequent step \(t=2,\dots,T\), the method performs \(K\) OT-based gradient steps. In each such step it computes the prior prediction
$$
\hat x_{t|t-1} \leftarrow \mathbf f(\hat x_{t-1}),
$$
$$
\Sigma_{t|t-1} \leftarrow \mathbf F_t\Sigma_{t-1}\mathbf F_t^\top + \hat{\mathbf Q}_{t-1},
$$
and the predictive measurement covariance
$$
\mathbf S_{t|t-1} = \mathbf H_t\Sigma_{t|t-1}\mathbf H_t^\top + \hat{\mathbf R}_{t-1}.
$$
It then samples source particles, computes the current innovation
$$
r_t = y_t - \mathbf h(\hat x_{t|t-1}),
$$
stores \(r_t\) in the residual queue, constructs the target support points
$$
\{\tilde y_t^j\}_{j=1}^W = \{y_t + r \mid r\in \mathrm{Res}\},
$$
forms the empirical measures, solves the OT problem, computes \(\mathcal L_{\mathrm{OT}}\), updates \(\theta\), and finally runs the adapted network forward:
$$
\hat x_t,\ \Sigma_{t|t},\ \mathbf Q_t,\ \mathbf R_t \leftarrow OTAKNet_\theta(y_t).
$$
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** [2508.07037].

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 Q\) and \(\mathbf R\) 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 [2508.07037].

## 5. Empirical evaluation on Lorenz and NCLT

The synthetic benchmark is a Lorenz attractor model with
$$
x_t = \exp\!\left( \begin{bmatrix} -10 & 10 & 0 \\ 28 & -1 & -x_{t-1,1} \\ 0 & x_{t-1,1} & -\frac83 \end{bmatrix} \Delta \right)x_{t-1} + \mathbf w_t,
$$
$$
y_t = \mathbf h(x_t)+\mathbf v_t,
$$
identity \(\mathbf h\), Gaussian noises, and \(\Delta=0.02\) [2508.07037]. The noise setting is
$$
\mathbf Q = q^2 \mathbf I,\qquad \mathbf R = r^2 \mathbf I,\qquad \nu \triangleq \frac{q^2}{r^2}.
$$
Performance is reported as MSE in dB across drift levels \(1/r^2 \in \{-10,0,10,20,30\}\) dB [2508.07037]. The baselines are Oracle EKF / OKF with true covariance values, VBAKF-PR, SWVAKF, AKNet, MAML-KalmanNet, KalmanNet, and OTAKNet [2508.07037].

The paper reports that OTAKNet consistently outperforms all non-oracle baselines across all tested Lorenz drift levels. Representative results include: at \(-10\) dB drift, KalmanNet \(15.26\), OTAKNet \(11.83\), SWVAKF \(15.11\), AKNet \(21.14\), and MAML-KalmanNet \(22.46\); at \(20\) dB drift, KalmanNet \(-16.35\), OTAKNet \(-18.78\), SWVAKF \(-13.46\), AKNet \(-13.10\), and MAML-KalmanNet \(-10.45\); and at \(30\) dB drift, KalmanNet \(-24.63\) versus OTAKNet \(-26.22\) [2508.07037]. The oracle EKF remains best because it uses the true noise covariances, but OTAKNet narrows that gap more effectively than the adaptive baselines considered [2508.07037]. Under \(1/r^2=20\) dB drift, OTAKNet is reported to adapt within about **25 time steps** and approach the performance of EKF with true covariances [2508.07037].

The real-world benchmark uses **NCLT** robot localization from IMU, odometry, and GPS. The state is
$$
x_t = [x_t,\ y_t,\ v_{x,t},\ v_{y,t},\ \theta_t]^\top,
$$
with dynamics
$$
x_t = \begin{bmatrix} x_{t-1} + \Delta v_{c,t}\cos(\theta) \\ y_{t-1} + \Delta v_{c,t}\sin(\theta) \\ v_{c,t}\cos(\theta) \\ v_{c,t}\sin(\theta) \\ \theta_t \end{bmatrix} + w_t,
$$
and measurement
$$
y_t = \begin{bmatrix} 1&0&0&0&0\\ 0&1&0&0&0 \end{bmatrix}x_t + v_t.
$$
The noise assumptions are
$$
\mathbf Q = \operatorname{diag}(1.0,1.0,0.001,0.001,0.001),
$$
$$
\mathbf R_{\mathrm{GPS}} = 10^2 \mathbf I_2.
$$
The trajectory is divided into 20 chunks of 200 steps: 13 train, 3 validation, and 4 test [2508.07037].

Two training regimes are reported: **Full training** with 13 labeled trajectories offline and **Limited training** with only 3 trajectories offline [2508.07037]. The quantitative results are as follows.

| Setting | Method | Result |
|---|---|---|
| Full training | EKF | \(7.51\pm2.55\) |
| Full training | VBAKF | \(7.12\pm2.42\) |
| Full training | SWVAKF | \(8.45\pm1.75\) |
| Full training | KalmanNet | \(6.35\pm2.15\) |
| Full training | OTAKNet | \(6.10\pm2.84\) |
| Limited training | KalmanNet | \(9.40\pm1.88\) |
| Limited training | OTAKNet | \(7.10\pm3.03\) |

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 [2508.07037]. The paper also notes that VBAKF outperforms fixed-noise EKF on NCLT, which suggests that real environmental drift is present in that dataset [2508.07037]. 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 [2508.07037].

## 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 [2508.07037]. On the synthetic benchmark the reported results are OTAKNet \(-18.45\pm0.41\), without warm-up \(-17.48\pm0.51\), and without OT \(-17.65\pm0.62\). On NCLT the reported results are OTAKNet \(6.10\pm2.84\), without warm-up \(6.28\pm2.54\), and without OT \(6.29\pm2.49\) [2508.07037]. 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
$$
\mu_{\mathrm{src}}(\theta) = \mathcal N\bigl(m_t(\theta), S_{t|t-1}(\theta)\bigr), \quad \mu_{\mathrm{tgt}} = \mathcal N\bigl(y_t,\widehat\Sigma_W\bigr),
$$
where
$$
m_t(\theta)=\mathbf h(\hat x_{t|t-1}(\theta)), \qquad S_{t|t-1}(\theta)=\mathbf H_t\Sigma_{t|t-1}(\theta)\mathbf H_t^\top+\mathbf R(\theta),
$$
and
$$
\ell_t(\theta)=W_2^2\bigl(\mu_{\mathrm{src}}(\theta),\mu_{\mathrm{tgt}}\bigr),
$$
then for Gaussians
$$
\ell_t(\theta) = \|m_t(\theta)-y_t\|^2 + \operatorname{Tr}\!\Bigl( S_{t|t-1}(\theta)+\Sigma_e -2\bigl(S_{t|t-1}(\theta)^{1/2}\Sigma_e S_{t|t-1}(\theta)^{1/2}\bigr)^{1/2} \Bigr).
$$
Under innovation stationarity and ergodicity, differentiability of the OT loss, and a Polyak–Łojasiewicz condition, gradient descent converges to a \(\theta^*\) such that
$$
S_{t|t-1}(\theta^*)=\Sigma_e.
$$
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 [2508.07037].

The computational cost is nontrivial. The paper states that Sinkhorn-type OT can be solved in \(O(n^2)\) per iteration and uses IPOT for practical differentiable computation [2508.07037]. Reported inference times in the Lorenz setting are EKF \(0.53\) s, SWVAKF \(4.28\) s, VBAKF \(19.81\) s, and OTAKNet \(21.02\) s [2508.07037]. 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 \(\mathbf f\) and \(\mathbf h\), relies on an **offline-trained KalmanNet-style backbone**, and uses a target-distribution construction \(y_t + e_{t-W+j}\) that is empirically motivated rather than uniquely derived [2508.07037]. It is not aimed at general non-Gaussian or multimodal filtering, despite the broader applicability of OT in other contexts [2508.07037].

## 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** [2201.12625]. 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 [2201.12625].

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 [2508.07037]. Similarly, it is not presented as merely updating \(\mathbf Q\) and \(\mathbf R\) through closed-form adaptive-KF formulas; the method updates the broader neural filter parameterization \(\theta\) [2508.07037].

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 [2508.07037]. 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.

Source: https://www.emergentmind.com/topics/otaknet