Adaptive Kalman-Informed Transformer (A-KIT)

• Adaptive Kalman-Informed Transformer (A-KIT) is a hybrid neural architecture that integrates classical Kalman filter recursions with transformer networks to enable robust state estimation.
• It encodes system dynamics, noise statistics, and observation history into tokens processed via self-attention, effectively replicating prediction and update steps.
• A-KIT offers practical improvements in online adaptation and performance, demonstrating near-optimal accuracy in sensor fusion, robotics, and complex dynamical environments.

An Adaptive Kalman-Informed Transformer (A-KIT) is a class of hybrid neural architectures that fuses the rigor of Kalman (and extended Kalman) filtering with the expressive capacity of Transformers to enable data-driven, robust state estimation and adaptation in dynamical systems. The key idea is to encode system dynamics, noise statistics, and observation history into a structured token sequence, processed by self-attention and learnable submodules, such that the network learns to approximate—up to arbitrary accuracy—the recursive estimation steps of classical (or adaptive) Kalman filters. A-KIT also generalizes to online adaptation of noise statistics, hyperparameters, and even model structure, facilitating resilient performance under non-stationarity, partial knowledge, or distribution shift.

1. Problem Formulation and Theoretical Foundations

A-KIT operates within the canonical finite-dimensional state-space model: $$x_{k+1} = F x_k + q_k,\quad q_k \sim \mathcal{N}(0, Q)$$

$$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$

where $F$ is the state transition matrix, $H_k$ is the (possibly time-varying) measurement matrix, and $Q$, $R$ are the process and measurement noise covariances. The Kalman filter recursions—explicitly implemented in A-KIT’s attention mechanism—are: $$\text{Prediction:}\quad x_{k|k-1} = F x_{k-1|k-1},\quad P_{k|k-1} = F P_{k-1|k-1} F^T + Q$$

$$\text{Update:}\quad K_k = P_{k|k-1} H_k^T (H_k P_{k|k-1} H_k^T + R)^{-1}$$

$$x_{k|k} = x_{k|k-1} + K_k (y_k - H_k x_{k|k-1}),\quad P_{k|k} = (I - K_k H_k) P_{k|k-1}$$

Transformers, when suitably structured, can in-context learn an approximation of these recursions, including the ability to adjust for time-varying noise and model parameters, and to infer missing parameters through implicit estimation (Akram et al., 2024, Goel et al., 2023). The theoretical equivalence is supported by mapping self-attention to a Nadaraya–Watson kernel smoother, which under high-temperature limits and appropriate embeddings, recovers the Kalman recursion up to an error that is uniform in time (Goel et al., 2023).

2. Transformer Architecture, Encoding, and Fusion with Kalman Filtering

A-KITs deploy a decoder-style Transformer, typically with L ≈ 16 layers, H ≈ 4 attention heads, and hidden dimension $d_\text{model} \approx 512$, leveraging GeLU activation and LayerNorm. The central innovation is in the construction of the input context: all system parameters, measurements, and structural tokens (e.g., "F-slot," "Q-slot") are flattened into a token sequence representing both observations ($$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$0), measurement matrices ($$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$1), and global parameters ($$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$2, $$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$3, $$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$4) (Akram et al., 2024). For time-varying or nonlinear settings, token slots also encode state-transition Jacobians (EKF), innovation-based covariance estimates, and context windows of inertial or sensor data (Cohen et al., 2024, Hohmeyer et al., 20 Nov 2025).

Within self-attention, certain heads specialize: some recover matrix multiplication (Mul) for state propagation, others implement affine updates, division, or apply the Kalman gain. This specialization is achieved through the learning/initialization of projection weights $$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$5, enabling the network to materialize the Kalman steps as a "Mul–Div–Affine–Update" sequence (i.e., realization of the complete Kalman filter algorithm in neural attention) (Akram et al., 2024).

3. Training Methodologies and Online Adaptation

Training involves curriculum learning over context length $$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$6, ramping noise scales $$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$7 over training steps, and minimizing losses on hidden state/MSE prediction: $$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$8

$$y_k = H_k x_k + r_k,\quad r_k \sim \mathcal{N}(0, R)$$9

Optimization is typically conducted using Adam (learning rate $F$0) with batch sizes around 64 (Akram et al., 2024). For nonlinear or time-varying noise, as in fusion settings or robotics, A-KIT learns a positive-definite diagonal scaling of the innovation-based $F$1, with the scale factors emitted by the Transformer, i.e., $F$2 (Cohen et al., 2024). Integrating such adaptive outputs within an EKF loop enables per-step online adaptation.

Online extension is readily achieved via sliding windowed context and auxiliary "parameter-estimation" submodules—small Transformer layers that specialize in estimating $F$3, etc., and feed discoveries back to the main state-estimation circuit (Akram et al., 2024, Goel et al., 2023). Bayesian variants propagate mean and covariance over the transformer's finetuning weights, performing weight updates via Kalman-style measurement correction to adapt the network itself online (Jing et al., 12 Sep 2025).

4. Empirical Evaluation and Robustness

Performance evaluation uses the Mean-Squared Prediction Difference (MSPD) and root mean square error (RMSE) compared to the classical Kalman filter and regression baselines: $F$4 A-KITs achieve MSE within 1% of the optimal Kalman filter on state estimation and one-step prediction tasks as $F$5 increases to 40 (for $F$6 latent states) (Akram et al., 2024). Withholding noise covariances ($F$7) from the context does not degrade performance: the Transformer implicitly infers these parameters, effectively emulating Dual Kalman filtering when $F$8 is also unknown. Extension to $F$9 outputs preserves this accuracy in the multivariate case.

In practical deployment (e.g., underwater vehicle navigation), A-KIT outperforms both classical EKF and model-based adaptive EKF baselines by average position RMSE improvements of ~49.5% and ~35.4%, respectively (Cohen et al., 2024). In humanoid robotics, a hybrid InEKF-Transformer design delivers sub-millimeter RMSE on position/orientation compared to conventional InEKF and outperforms RNN-based KalmanNet by orders of magnitude (Hohmeyer et al., 20 Nov 2025).

5. Interpretability, Limitations, and Open Directions

Interpretable analysis demonstrates that specific attention heads implement the normalized Kalman gain

$H_k$0

and realize the covariance update via affine and division operations on token slots (Akram et al., 2024). Adaptive heads, gating, and context windows enable robust tracking under parameter drift or abrupt changes in dynamics.

Key limitations include:

• Scope largely restricted to white (i.i.d.) noise models; extension to colored noise or temporally correlated process noise is non-trivial.
• Scalability concerns: large $H_k$1 or $H_k$2 entails quadratic memory in the number of tokens, motivating the need for low-rank or sparse attention implementations.
• Nonlinear models (i.e., extended or unscented Kalman filtering) require integrating nonlinear modules and possibly structured positional encodings.
• Closed-form theoretical guarantees on convergence under arbitrary parameter drift or in settings with complex, partially observed dynamics remain largely open (Akram et al., 2024, Goel et al., 2023).

Challenges in training include mitigating "exposure bias" in autoregressive setups, handling multi-rate sensor data streams, and avoiding over-correction under sampling jitter—especially salient for high-dimensional robotic control (Hohmeyer et al., 20 Nov 2025). Sim-to-real transfer is best handled by joint training on simulated and real data, introducing domain randomization, and explicitly modeling timestep variations.

6. Generalizations and Applications

The A-KIT paradigm is broadly applicable to sensor-fusion (IMU+GNSS, inertial navigation, visual odometry), adaptive control (LQG) (Goel et al., 2023), and online sequential learning under severe memory constraints (Jing et al., 12 Sep 2025). End-to-end differentiable architectures permit direct minimization of state-estimation error with supervised, reinforcement, or likelihood-based objectives. Measurement-feedback controllers can be folded into the A-KIT structure, training "control heads" alongside state estimation while maintaining closed-loop stability guarantees up to uniform time error if the underlying self-attention sufficiently approximates Kalman operations (Goel et al., 2023).

A summary of A-KIT properties is shown below:

Property	Classical Kalman Filter	A-KIT (Transformer-based)
Model Structure Required	Full (F, Q, R, H)	Can operate with missing/partial
Online Adaptation	Hand-tuned/adaptive AEKF	Unified via transformer attention and gating
Nonlinear/Multirate	EKF/UKF extensions	Patch-embedding, nonlinear MLPs
Empirical Performance	Optimal (MSE)	Matches or exceeds, adapts online
Scalability	Linear in state dim	Quadratic in # tokens, mitigated by efficient attention

Transformers thus provide a flexible, data-efficient, and theoretically principled foundation for modern state estimation, merging the statistical guarantees of Kalman filtering with the data-adaptive inference and representation learning of deep networks (Akram et al., 2024, Goel et al., 2023, Cohen et al., 2024, Jing et al., 12 Sep 2025, Hohmeyer et al., 20 Nov 2025).