Nonlinear Kalman Filtering based on Self-Attention Mechanism and Lattice Trajectory Piecewise Linear Approximation

This section introduces the operation of the attention mechanism and our proposed simplified attention network. The self-attention mechanism is shown in Fig. 1(a). It transforms the input sequence ${x_{1},x_{2},\ldots,x_{n}}$ into query, key, and value matrices through linear mappings. For each sequence element $x_{i}$, it calculates the dot products with all sequence elements, forming an attention distribution through softmax that indicates the elements’ dependencies. This attention distribution is then multiplied with value to produce the output sequence $\{v_{1},v_{2},\ldots,v_{n}\}$, where each element $v_{i}$ integrates information from the entire sequence. This ensures each processed element reflects the influence of every other element, overcoming the limitations of distance.

A simplified version of this mechanism, as depicted in Fig. 1(b), streamlines the process by using a single matrix $X$ for multiplying both the attention distribution and the output sequence. This is feasible because $d_{embed}=d_{model}$, and the dimension of $X$ matches that of the input sequence matrix, eliminating the need for separate and extra linear mappings. Given the small amount of sequence data and the simple distribution of features in the Kalman filtering process, employing the full multi-head attention mechanism can increase training difficulty and lead to overfitting. By reducing the number of parameters, this simplified approach boosts efficiency without sacrificing the model’s ability to capture crucial dependencies.

The LTPWL expression is an approximate method of constructing a lattice PWL expression for a nonlinear function. It can simultaneously accomplish the linearization of the nonlinear system and the construction of lattice PWL expression. Its main process involves selecting a set of linearization points along the system’s state trajectory, constructing linear segments at each of these points, and finally using these segments to build a lattice PWL expression to approximate the original nonlinear system.

Taking the nonlinear function shown in Fig. 2 as an example, by selecting points $x_{1},x_{2},$ and $x_{3}$ along the state trajectory (here, the $x$-axis), and creating corresponding linear segments $l_{1},l_{2},$ and $l_{3}$ at each point, we represent the PWL function as shown in (1),

this method simplifies the approximation of nonlinear functions without defining specific interval ranges for $x$, making it a more compact solution compared with traditional piecewise form (2),

The general form of the lattice PWL is shown in (3),

where $N$ is the number of base regions $\mathbb{D}_{i}$, with each $\mathbb{D}_{i}$ being a polyhedron satisfying $\{x|l_{i}(x)=l_{j}(x),j\neq i\}\cap\mathbb{D}_{i}=\emptyset$. Furthermore, $\tilde{I}_{\geq,i}$ denotes the index set of affine functions in the base region $\mathbb{D}_{i}$ that are greater than or equal to the linear segment $l_{i}(x)$, i.e., $\tilde{I}_{\geq,i}=\{j|l_{j}(x)\geq l_{i}(x),\forall x\in\mathbb{D}_{i}\}$.

For the PWL function shown in Fig. 2 with three basic regions, $N=3$. Taking the linear segment $l_{1}$ as an example, in the basic region $\mathbb{D}_{1}$, the affine functions that are greater than or equal to $l_{1}$ are $l_{1}$ and $l_{2}$. Then, a minimum operation is applied to $l_{1}$ and $l_{2}$ to construct a “term” in the lattice PWL expression, resulting in $\mathop{\min}\{l_{1},l_{2}\}$. Terms for linear segments $l_{2}$ and $l_{3}$ are constructed in a similar way, resulting in $\mathop{\min}\{l_{1},l_{2}\}$ and $\mathop{\min}\{l_{1},l_{3}\}$, respectively. Then, a maximum operation on these three terms yields the final lattice PWL expression (1). When evaluating the lattice piecewise linear expression, it is only necessary to substitute the value of the variable $x$, without the need to consider the interval range of the variable as in the traditional piecewise form. Moreover, we can conveniently identify the linear segment $l_{i}$ that represents the current system dynamics through comparison operations. As we will see in Section III-D2, this property of the lattice expression is particularly suited for batch estimation algorithms to generate pre-training data.

Considering the discrete-time nonlinear system given by (4),

where $f$ and $h$ represent the nonlinear state transition and observation functions, respectively, $x_{k}$ denotes the state vector at time step $k$, and $y_{k}$ represents the observation at $k$. $w_{k}$ and $v_{k}$ correspond to the process noise and observation noise at $k$, respectively, both assumed to be Gaussian white noise with their covariance matrices $Q$ and $R$.

As shown in Fig. 3(a), the framework aligns with the traditional Kalman filtering algorithm, using the system model for state recursion and output prediction at each time step. At any given moment, such as time step $k$, $\check{x}_{k}$ and $\hat{y}_{k}$ represent the prior estimate of the state $x_{k}$ and the predicted system output, respectively, while $\hat{x}_{k}$ is the posterior state estimate. A self-attention mechanism is incorporated to predict the Kalman gain $K_{k}$. The gain, acting as a fusion coefficient between model predictions and system observations, adjusts based on process $w_{k}$ and observation noise $v_{k}$. By leveraging the self-attention mechanism, the network captures sequential data dependencies more effectively, leading to enhanced estimation accuracy by fitting the system’s noise characteristics.

The input to the self-attention mechanism network comprises two types of features. which can be expressed as in (5) at time step $k$,

$\Delta x_{k-1}\in\mathbb{R}^{m}$ represents the forward update difference, and $\Delta y_{k}\in\mathbb{R}^{n}$ represents the innovation. To fully utilize the sequence processing capability of the self-attention network and to capture the features among state sequences thoroughly, values from a past time window for each feature are collected to form a sequence, which then serves as the network input. Moreover, as the filtering process goes, the window slides accordingly. Assuming the size of the time window is $s$, the network input at $k$ can be represented as in (6),

If the historical data is insufficient (i.e., when $k<s$), the network fills missing feature values in the input sequence with zeros, as shown in Fig. 3(a). Then the network uses the processed input sequence to predict the Kalman gain $K_{k}$, which is subsequently utilized to update the system state.

Similar to the traditional KF, the overall filtering framework can also be summarized into prediction and update steps:

• •

  Prediction step:

  	$\displaystyle\check{x}_{k}$	$\displaystyle=f(\hat{x}_{k-1}),$		(7a)
  	$\displaystyle\hat{y}_{k}$	$\displaystyle=h(\check{x}_{k}).$		(7b)

• •

  Update step:

  	$\displaystyle K_{k}$	$\displaystyle=\text{attention}(\Delta\hat{x}_{k-1},\Delta y_{k}),$		(8a)
  	$\displaystyle\hat{x}_{k}$	$\displaystyle=\check{x}_{k}+K_{k}(y_{k}-\hat{y}_{k}).$		(8b)

As shown in Fig. 3(b), the entire self-attention network consists of two linear embedding layers for initial processing, a positional encoding layer to integrate sequence position information, a simplified attention layer for capturing dependencies within the sequence, and two fully connected layers within a multi-layer perceptron (MLP) block to enrich feature representation. It concludes with a linear mapping layer that projects the processed features to predict the Kalman gain $K_{k}$. Starting with input features $\Delta x_{k-1}\in\mathbb{R}^{(B,s,m)}$ and $\Delta y_{k}\in\mathbb{R}^{(B,s,n)}$, the network fuses and transforms these inputs into a sequence $X_{\text{input}}\in\mathbb{R}^{(B,2s,d_{\text{model}})}$ via the embedding layers. $X$ is then enhanced for dependency recognition in the attention layer and further processed by the MLP to capture non-linear characteristics,  finally outputting the predicted Kalman gain $K_{k}\in\mathbb{R}^{(B,m*n)}$.

The system function is given by (19), where both the system state and output are two-dimensional vectors, i.e., $x,y\in\mathbb{R}^{2}$. The parameters of the system are shown in Table I.

${Para}_{s}$ represents the true parameters of the system, and ${Para}_{m}$ denotes the parameters of the model. The system’s state transition function (19a) and observation function (19b) are both nonlinear. $w_{k}$ and $v_{k}$ represent the process noise and observation noise, respectively, both assumed to be Gaussian white noise with covariance matrices denoted by $Q$ and $R$.

We use the original nonlinear model (19), starting with the initial state $x_{0}=[0.1,0.1]^{\rm{T}}$. datasets for training, validation, and testing are generated under random noise. The training dataset contains $N=1000$ data entries, each with $L=10$ time steps. The validation dataset contains $N=100$ data entries, each with $L=10$ time steps. The test dataset contains $N=200$ data entries, each with $L=100$ time steps. Moreover, to generate pre-training data, a noise-free state trajectory of 10 time steps is produced using the same nonlinear model and initial state.

For training parameters, the self-attention network’s sliding window size was set to $s=4$, with a batch size of 50 and a learning rate of 1e-4. For the AtKF, pre-training was conducted for 50 epochs and the subsequent training phase for 20 epochs. To ensure fairness, KalmanNet was trained for 70 epochs. All training processes are conducted on a GTX-3090 GPU.