Title: Gated Inference Network: Inference and Learning State-Space Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: Time Series and Recurrent Networks
Abstract: This paper advances temporal reasoning within dynamically changing high-dimensional noisy observations, focusing on a latent space that characterizes the nonlinear dynamics of objects in their environment. We introduce the *Gated Inference Network* (GIN), an efficient approximate Bayesian inference algorithm for state space models (SSMs) with nonlinear state transitions and emissions. GIN disentangles two latent representations: one representing the object derived from a nonlinear mapping model, and another representing the latent state describing its dynamics. This disentanglement enables direct state estimation and missing data imputation as the world evolves. To infer the latent state, we utilize a deep extended Kalman filter (EKF) approach that integrates a novel compact RNN structure to compute both the Kalman Gain (KG) and smoothing gain (SG), completing the data flow. This design results in a computational cost per step that is linearly faster than EKF but introduces issues such as the exploding gradient problem. To mitigate the exploding gradients caused by the compact RNN structure in our model, we propose a specialized learning method that ensures stable training and inference. The model is then trained end-to-end on videos depicting a diverse range of simulated and real-world physical systems, and outperforms its ounterparts —RNNs, autoregressive models, and variational approaches— in state estimation and missing data imputation tasks.

Gated Inference Network: Inference and Learning State-Space Models

Abstract

This paper advances temporal reasoning within dynamically changing highdimensional noisy observations, focusing on a latent space that characterizes the nonlinear dynamics of objects in their environment. We introduce the Gated Inference Network (GIN), an efficient approximate Bayesian inference algorithm for state space models (SSMs) with nonlinear state transitions and emissions. GIN disentangles two latent representations: one representing the object derived from a nonlinear mapping model, and another representing the latent state describing its dynamics. This disentanglement enables direct state estimation and missing data imputation as the world evolves. To infer the latent state, we utilize a deep extended Kalman filter (EKF) approach that integrates a novel compact RNN structure to compute both the Kalman Gain (KG) and smoothing gain (SG), completing the data flow. This design results in a computational cost per step that is linearly faster than EKF but introduces issues such as the exploding gradient problem. To mitigate the exploding gradients caused by the compact RNN structure in our model, we propose a specialized learning method that ensures stable training and inference. The model is then trained end-to-end on videos depicting a diverse range of simulated and real-world physical systems, and outperforms its counterparts -RNNs, autoregressive models, and variational approaches- in state estimation and missing data imputation tasks.

1 Introduction

A state-space model is a type of graphical model that effectively represents noise-afflicted data [1]. Typically, the objective is to conduct inference, which involves obtaining accurate estimates of the posterior distribution of latent states or noise-free measurements. In videos, inferring the state space presents a complex challenge owing to the high dimensionality of noisy observations, which only offer partial information about the states. Previous methods have proposed various inference approaches such as sampling [2], variational inference [3], or belief propagation [4]. Within the framework of a hidden Markov process, classical methods like the celebrated Kalman filter (KF) and smoother [5, 6], along with their extensions like EKF and UKF, [7, 8] have been employed to address the posterior inference problem. However, the reliance of these methods on accurate knowledge of the underlying dynamics, coupled with computationally intensive matrix inversions, poses challenges in scalability to high-dimensional latent spaces.

To address these issues, this paper distinguishes between the sensory observation and its dynamics by disentangling two representations: a transformed observation w t obtained by mapping the original sensory observation o t through a nonlinear function modeled by neural networks, and the latent state space x t describing the dynamics of w t at time step t . To infer high dimensional state space (dynamics) x t , we introduce filtering and smoothing techniques that build upon classical EKF updates but incorporate flexible function estimators without relying on a constrained graphical model. We model temporal dynamics with a linearized Gaussian SSM, which is adapted to accommodate complex dynamics. To address non-linearity and the multiplicity of dynamics, we learn multiple transition and emission matrices, each modeling a unique dynamical scenario. Then, the weight of each matrix during inference is determined by the past state x t -1 through a neural network. To prevent mode collapse and ensure the system's ability to capture diverse dynamics, we introduce a loss term proportional to the KL divergence of the transition matrices during training. This design enables our model to handle various state scenarios, similar to approaches found in the switching linear dynamics systems (SLDS) literature [9]. Unlike EKF, which requires precise characterization and modeling of the underlying dynamics, we operate under the assumption that noise statistics and the underlying physical dynamics of SSM are entirely unknown.

The main drawback of the EKF is that it takes O ( n 3 ) time per step, because we need to invert the posterior covariance matrix with size of n . This makes the method slow to use. Moreover, we identify the computation of the KG and SG in the (extended) KF as critical components that rely on noise statistics and domain knowledge. In this paper, we efficiently model KG and SG using GRU cells, bypassing the need for matrix inversion in the KF flow. Consequently, the GIN algorithm achieves a time complexity of O ( β 2 n 2 ) per step, making it linearly faster than EKF. To capitalize on the sparsity of the covariance matrix, we employ a convolutional approach that further reduces the size of n relative to β , where 0 < β < 1 . However, the nonlinear structure of GRU cells can lead to gradient explosion [10], a phenomenon that occurs when the dynamics undergo significant changes as GRU parameters traverse bifurcation points [11] during the learning process. To address this issue, we propose a learning method based on an analysis of GRU cell dynamics. This method aims to prevent parameter values from reaching bifurcation points by stabilizing the dynamics of GRU cells, thereby addressing the issue of gradient explosion.

To summarize, our primary contribution is a novel algorithm designed for the recursive inference of SSMs with arbitrary nonlinear (and possibly non-Gaussian) emission and transition models, with high-dimensional noisy stimuli observation. This algorithm utilizes a linearized Gaussian dynamics approach, with parameter estimation conducted through neural networks. This means that we can apply the recursive Bayesian updates, akin to the Kalman filter and Kalman smoother. We introduce a likelihood for inferring the states (dynamics) and another likelihood for inferring the high-dimensional images (of video). The objective is maximized in a supervised manner, depending on the task. We propose learning schemes for GRU cells to address issues related to gradient explosion and instability. Finally, we introduce a loss term proportional to the KL divergence of the learned transition matrices to prevent the system from becoming stuck in mode collapse.

To substantiate our claims, we conduct five experiments. First, we simulate a nonlinear dynamic system using the pendulum sequence video, a common benchmark in this literature, to demonstrate our model's ability to infer both dynamics and images. Second, we introduce a more challenging experiment of a simulated nonlinear double pendulum, where the video sequence is heavily distorted with noise, showcasing our model's resilience of dynamics estimation and image imputation to noisy observations. Third, we present a switching dynamics irregular bouncing ball experiment to illustrate our model's ability to handle multiple dynamic scenarios. Fourth, we perform visual odometry using the KITTI dataset, demonstrating the practical applicability of our methods in real-world scenarios. Finally, we assess the effectiveness of our proposed gradient explosion handling scheme by evaluating convergence across various simulation seeds.

2 Related Works

EKF and UKF represent early extensions of the original Kalman filter, allowing for nonlinear transitions and emissions. These methods, categorized as model-based (MB) algorithms, rely on precise knowledge and modeling of the underlying dynamics within a fully characterized SSM. Consequently, the performance of these MB methods is heavily contingent on the accuracy of domain knowledge and model assumptions. Recent approaches, such as BackpropKF [12] and SIN [13], perform EKF in the latent state using deep encoders for feature extraction. However, they face similar

limitations as EKF, including computational complexity and scalability. To address these issues, a common approach involves considering a diagonal covariance matrix in the KF flow as in [14] and [15]. However, we aim to find a richer approximation to the posterior.

A significant body of work, including Variational Auto-Encoders (VAE) [3], Embed to Control (E2C) [16], and Importance Weighted VAE (IWVAE) [17], integrates deep learning with variational inference (VI), but lacks memory cells and recurrent structures for handling imputation tasks. To address this, EM-based VI approaches like Structure VAE (SVAE) [18], Kalman VAE (KVAE) [19], Disentangled VAE (DVAE) [20], Extended KVAE (EKVAE) [21], Robust VAE [22], and Markovian VAE (MVAE) [23] use original KF equations for filtering and smoothing. However, they cannot directly optimize states (dynamics), as noted in RKN [24] and CRU [25]. Classical memory networks like LSTMs [26], GRUs [27], and simple RNNs [28] infer latent states but fail to provide insights into uncertainties and dynamics.

LatentNet [29], KalmanNet [30] and SSI-SSM [31] utilize GRU in their structures for state updates, similar to the parameterization used in the GIN. However, these models typically necessitate access to complete or partial dynamics information, i.e. they are MB, and their reliance on vanilla GRU cells can cause instability of the SSM. Liu et al. [32] proposes a disentanglement model for audio data, which shares similarities with the GIN. However, their parameterization and objectives differ from those of the GIN, and they do not utilize compact RNNs in their inference structure.

We cover the most relevant works in this section. A more detailed discussion, including a comprehensive comparison of recent related works like System Identification (SI) with details of EM algorithm, auto-regressive (AR) and SLDS models, is provided in Appendix A.8.3 and Table 7. This builds upon the findings of the RKN with additional complements. We also conduct an empirical complexity analysis in appendix A.8.2, to evaluate the computational efficiency of our method compared to the discussed approaches. This analysis involved measuring the execution time per iteration using a clock on the wall as the benchmark.

3 Background

Linear Gaussian state space models. Linear Gaussian state space models (LGSSMs) are commonly employed to model vectors of observed time series, denoted as w = [ w 1 , ..., w T ] . LGSSMs excel in filtering and smoothing tasks. They model the first-order Markov process on the state space x = [ x 1 , ..., x T ] , which may also incorporate external control inputs u = [ u 1 , ..., u T ] under the multivariate normality assumption of the state:

γ t represents the parameters of the system at time t , encapsulating information from matrices F t , B t , H t , Q t and R t , which correspond to the state transition, control, emission, process noise covariance, and observation noise covariance, respectively. At each time step, the transition and emission procedures are subject to distortion caused by process noise and observation noise, following distributions of N ( 0 , Q t ) and N ( 0 , R t ) , respectively.

Filtering and smoothing algorithms. The goal of filtering is to infer the posterior distribution p γ 1: t ( x t | w 1: t , u 1: t ) , while, the smoothing aims to infer the smoothing distribution p γ 1: T ( x t | w 1: T , u 1: T ) given the whole observations. Without loss of generality, we drop the input variable u in the rest of paper, while one can condition all of the distributions on the input variable in the case of its existence. Considering the prior state parameterization p γ 1: t ( x t | w 1: t -1 ) = N ( µ t | t -1 , Σ t | t -1 ) and relying on [33], the filtering and smoothing parameterizations are as:

𝑞𝑡-1

x𝑡-1

w𝑡-1

𝑞𝑡

x𝑡

w𝑡

𝑞𝑡+1

x𝑡+1

w𝑡+1

…

…

…

…

…

…

o𝑡-1

s

𝑡-1

o𝑡

s

𝑡

o𝑡+1

s

𝑡+1

with S f = [ H t Σ t | t -1 H T t + R t ] and S s = [ Σ t +1 | T -Σ t +1 | t ] . In (4), K t and J t are KG and SG, respectively. See appendix A.1 for a full derivation of filtering-smoothing distributions.

4 Gated Inference Network

The background section defines key elements for the GIN: original observations o 1: T , task-dependent output which can be either the original denoised-imputed observations o 1: T or physical system's states s 1: T , transferred observations w 1: T ∈ R m × T , latent states (dynamics) x 1: T ∈ R n × T , noise covariance of transferred observations R 1: T with diagonals r 1: T ∈ R m × T , noise covariance of states process Q 1: T with diagonals q 1: T ∈ R n × T , and the parameters of the SSM γ 1: T . We consider the dimensions of the transferred observation and state at each time step to be m and n , respectively. The states evolve through transition distribution p γ 1: T ( x 1: T ) = p ( x 1 ) ∏ T t =2 p γ t ( x t | x t -1 ) . Each w t is assumed to be drawn from noisy emission probability p γ t ( w t | x t ) , then the generative process is assumed to factorize as p γ 1: T ( x 1: T , w 1: T ) = p γ 1: T ( x 1: T ) ∏ T t =1 p γ t ( w t | x t ) . This factorisation imposes emission conditional independence like w t ⊥ ⊥ ( w -t , x -t ) | x t , where ( w -t , x -t ) = ( w (1: t -1 ,t +1: T ) , x (1: t -1 ,t +1: T ) ) . Common models that assume conditional independence include linear dynamical systems, hidden Markov models, but also nonlinear SSMs with higher-order Markov chains in latent space. Graphical model of the GIN is in figure 2, where output (one of colored nodes) is generated, contingent upon the task. In all of our scenarios, the parameters of the SSM, γ 1: T , are completely unknown.

From SI perspective, the GIN operates similarly to a Hammerstein-Wiener (HW) model [34], employing non-linear transfer functions e ( . ) and d ( . ) (Figure 3). Leveraging encoder-decoder structures for e ( . ) and d ( . ) , the GIN conducts the transfer between original sensory observations o 1: T and lowerdimensional representations w 1: T and task dependent output. The transition block in Figure 3 evolves states using the proposed deep EKF approach, efficiently approximating filtering-smoothing distributions while ensuring system stability with the imposed constraint. Further details and formulations are provided in the inference section. The transition block depicted in Figure 4.

5 Parameterization and Inference

Given original observations o 1: T and transferred observations w 1: T , our aim is to infer the latent states x 1: T . To accomplish this, we seek to infer the marginal distributions

𝑒(. )

Transition Block

𝑑(. )

𝒐1:𝑇

(𝒘1:𝑇 , 𝒓

1:𝑇

)

( 𝒙

𝟏:𝑇

|𝒘1:𝑇 )

{𝑠

1:𝑇,

𝑜 1:𝑇

}

Dynamic Net

Filtering

x t-1

|w1:𝑡-1

෠

F𝑡

෡

H𝑡

Smoothing

x t

|w1:𝑡

w𝑡

x t

|w1:𝑇

x t-1

|w1:𝑇

Dynamic Net

Filtering

x t

|w1:𝑡

෠

F𝑡+1

෡

H𝑡+1

Smoothing

x t+1

|w1:𝑡+1

w𝑡+1

x t+1

|w1:𝑇

x t

|w1:𝑇

Dynamic Net

Filtering

x t-2

|w1:𝑡-2

෠

F𝑡-1

෡

H𝑡-1

Smoothing

x t-1

|w1:𝑡-1

w𝑡-1

x t-1

|w1:𝑇

x t-2

|w1:𝑇

…

…

p γ 1: t ( x t | w 1: t ) for the online inference approach (filtering) and p γ 1: T ( x t | w 1: T ) for the full inference approach (smoothing). We introduce an advantageous prediction parameterization as p γ t ( x t | x t -1 , w 1: t -1 ) = N ( F t x t -1 , Q t ) , where x t -1 is sampled from the last filtering distribution, p γ 1: t -1 ( x t -1 | w 1: t -1 ) . Then, we obtain the prior distribution at time t , p γ 1: t ( x t | w 1: t -1 ) = N ( F t µ t -1 | t -1 , F t Σ t -1 | t -1 F T t + Q t ) = N ( µ t | t -1 , Σ t | t -1 ) , by marginalizing out x t -1 from ∫ p γ t ( x t | x t -1 , w 1: t -1 ) p γ 1: t -1 ( x t -1 | w 1: t -1 ) d x t -1 integration. The Gaussianity of p γ 1: t ( x t | w 1: t -1 ) results from the Gaussianity of prediction parameterization. After obtaining p γ 1: t ( x t | w 1: t -1 ) and observing w t , it is feasible to derive the filtering parameterization using equation (2). Once all transferred observations w 1: T are available, backward induction can be employed to propagate to previous states using the chain rule. This procedure, known as smoothing, is parameterized with (3). These parameterizations provide valuable insights into two key aspects: 1) appropriately modeling γ using neural networks, and 2) illustrating a tractable method to parameterize KG and SG and construct distributions approximations. These approximations serve the basis for constructing output.

Learning γ . To handle multiple dynamic scenarios, we learn K sets of state transition and emission matrices ˇ F k and ˇ H k , each representing a distinct dynamic status. These matrices are combined with state-dependent coefficients α k ( x t -1 ) as :

We model α k ( x t -1 ) with a K dimension softmax output neural network called the Dynamic Net . It takes the last filtering state x t -1 as input, containing the system's history with lower noise distortion than the transferred observations w t -1 . This choice enhances noise robustness, demonstrated in our experiments with time-correlated noise (See Appendix A.6).

In the graphical model shown in Figure 2, we observe two paths for belief propagation from x t -1 to x t . The first path, x t -1 → x t , linked with ˆ F t . The second path involves an intermediate variable q t ∼ N ( 0 , Q t ) : x t -1 → q t → x t . Since learned ˆ F t transfers information from x t -1 to x t , we argue that it can capture effects similar to those of Q t , as both are intended to convey message between x t -1 and x t [1]. By incorporating ˆ F t from equation (5) into the prior state distribution described earlier as p γ 1: t ( x t | w 1: t -1 ) = N ( F t µ t -1 | t -1 , F t Σ t -1 | t -1 F T t + Q t ) , we can neglect the Q t term in the covariance, as its influence is accounted for by the learned ˆ F t . There are two other meaningful parameterizations for the process noise matrix. First involves direct parameterization by x t -1 using a neural network. Second approach is to parameterize it recursively using x t -1 and q t -1 , resulting in a new graphical model with an edge from q t -1 → q t (See Appendix A.2). Both parameterizations are included in our results, presented in the appendix, to show their effectiveness.

Finally, the diagonal elements of the transferred observation noise vector r t are directly mapped from the original observation space using the encoder function e ( . ) shown in Figure 3. The mapping employs an activation function, elu+1, to handle the positivity of the diagonal elements.

Filtering and Smoothing Approximation. We consider GRU KG network to approximate KG as:

In Equation (6), Conv ( . ) represents a zero bias convolutional layer with pooling, employed to deal with sparsity of covariance matrix and extract its relevant information while reducing its size. The [ , ] symbol denotes the concatenation operator. Furthermore, the presence of a positive r t vector and the consideration of the Cholesky factor, M t M T t in (6), ensure the resulting covariance matrices maintain positive definiteness. This parameterization construct a new filtering distribution q γ 1: t ( x t | w 1: t ) = N (ˆ µ t | t , ˆ Σ t | t ) that is an approximation of (2). Consequently, we consider the approximated prior distribution as q γ 1: t ( x t | w 1: t -1 ) = N ( ˆ F t ˆ µ t -1 | t -1 , ˆ F t ˆ Σ t -1 | t -1 ˆ F T t ) = N (ˆ µ t | t -1 , ˆ Σ t | t -1 ) , where (ˆ µ t | t -1 , ˆ Σ t | t -1 ) are used in (7).

After obtaining filtering states from q γ 1: t ( x t | w 1: t ) , we use GRU SG network to approximate SG in a similar way we used GRU KG in (6) as:

where the first smoothing state is set to the last filtering state. The new smoothing distribution q γ 1: T ( x t | w 1: T ) = N (ˆ µ t | T , ˆ Σ t | T ) is an approximation of the exact smoothing distribution in (3). For a GRU cell with an input size of i and a hidden state size of h , the computational complexity in the forward pass is O (3 h ( h + i +3)) , which scales linearly with the input size [35]. Considering Conv ( Σ ) ∈ R β 2 n 2 as the input of our GRU cells in (6) and (8) with β ∈ [0 , 1] as the pooling ratio, the forward pass of GIN for one time step has a time complexity of O (3 hβ 2 n 2 ) , where n ≫ h . Compared with LGSSM matrix inversion time complexity of O ( n 3 ) , GIN is faster by a factor of n 3 hβ 2 , which is crucial in high-dimensional regimes.

6 Fitting

In the state estimation task, the output from d ( . ) in figure 3 equals s , the physical system's states. However, in the imputation task, output of d ( . ) is same as o since the original observation is reconstructed (see decoder structure for each task in A.8.1). The conditional distributions p ( w 1: T | o 1: T ) and q γ 1: T ( x 1: T | w 1: T ) are modeled using an encoder e ( . ) and smoothing parameterization, respectively. Meanwhile, the conditional distributions p ( s 1: T | x 1: T , w 1: T , o 1: T ) and p ( o 1: T | x 1: T , w 1: T ) are represented by the decoder d ( . ) for the tasks of state estimation and imputation. These distributions are modeled using multivariate Gaussian and Bernoulli distributions, respectively. Depending on the characteristics of the observations and states, alternative likelihood distributions can also be employed. For instance, the beta likelihood for data in the unit interval, mixtures for multiple marginal distributions, and the negative-binomial likelihood for positive count data (See appendix A.3).

Likelihood for Inferring States. The following theorem defines objective for inferring physical system's states. The proof is provided in appendix A.4.1.

Theorem 1. The lower bounded conditional log likelihood of the physical system's states given original images is determined as:

where the d m ( . ) and d c ( . ) determines those parts of the decoder obtaining the state mean and variance. N sequences of ( x ( i ) 1: T , w ( i ) 1: T ) ∼ q γ 1: T ( x 1: T , w 1: T | o 1: T ) are sampled for Monte Carlo integration estimation.

Likelihood for Inferring Images. For the imputation task, consider the ground truth as the sequence of images o 1: T with the dimension of D o . The following theorem defines objective for inferring images. See appendix A.4.2 for the proof.

Theorem 2. The lower bound of log likelihood of the original images is:

with N sequences of ( x ( i ) 1: T , w ( i ) 1: T ) ∼ q γ 1: T ( x 1: T , w 1: T | o 1: T ) for Monte Carlo integration estimation. d k ( . ) defines the corresponding part of the decoder that maps the k -th pixel of o t image.

Weuse Wishart distribution as a prior for our covariance matrix of states in (10) and (11), which pushes the covariance toward a scale of identity matrix. Such prior prevents getting high log-likelihood due to the high uncertainty. We shrink this scale toward zero as time passes, as we expect the model to finally perform with very little uncertainty, approaching deterministically.

Mode collapse handling. To address the issue of mode collapse, where the model becomes stuck in the same state, we propose a loss term on ( i, j ) pair of transition matrices as follows:

̸

with m i and m j being distinct hyper parameter vectors, and || . || F is the Frobenius norm. The term -1 2 ∑ K i =1 ∑ K j =1 ,j = i l mc ( i, j ) can be added into (10) and (11). This addition ensures that each ˇ F i captures a unique transition, distinct from others. From a statistical perspective, (12) represents a term proportional to -KL ( p i ( ˇ F ) || p j ( ˇ F ) .pri j ( ˇ F ) ) + KL ( p i ( ˇ F ) || pri i ( ˇ F ) ) (See appendix A.5). p i ( ˇ F ) ∼ MN n × n ( ˇ F i , I , I ) and p j ( ˇ F ) ∼ MN n × n ( ˇ F j , I , I ) , are matrix normal distributions with priors pri i ( ˇ F ) ∼ MN n × n ( Diag ( m i ) , I , I ) and pri j ( ˇ F ) ∼ MN n × n ( Diag ( m j ) , I , I ) , respectively.

Gradient explosion handling. Training (10) and (11) with SGD can be disrupted by a gradient explosion. Specifically, we limit the optimization in (10) and (11) subject to σ 1 ( U h ) < 2 , where σ i ( . ) represents the i -th largest singular value and U h ∈ R h × h denotes the weight matrix of the hidden state h ∈ R h with h dimension in the GRU cell (see eq.(5) in [35] or (53) in appendix A.4.3).

Theorem 3. When σ 1 ( U h ) < 2 , a GRU cell is locally stable at a fix point h ∗ = 0 .

We use Lyapunov stability and two lemmas for the proof of Theorem 3 in appendix A.4.3. Solving (10) and (11) with SGD gives an updated ˆ U h in each iteration that may not satisfy σ 1 ( ˆ U h ) < 2 . We offer two solutions to satisfy the constraint outlined in theorem 3. The first solution relies on the spectral theorem, offering higher accuracy but with increased computational cost. Alternatively, the second solution employs the Gershgorin circle theorem, providing lower accuracy but at a reduced cost. Due to space limitations, we provide the details of the second solution in appendix A.4.5. For the first solution, we modify ˆ U h in three steps: (i) Decompose ˆ U h by singular value decomposition (SVD) such that ˆ U h = WΛV . (ii) Replace the singular values of diagonal Λ that are greater than 2 with the threshold 2 -δ and obtain Λ = Diag ( min ( σ 1 , 2 -δ ) , ..., min ( σ h , 2 -δ )) . (iii) Reconstruct U h ← WΛV . Employing SVD approach, the cost of this solution is proportional to O ( h 3 ) . In the next theorem we show that constructed U h is the optimum solution. The proof is in appendix A.4.4.

Theorem 4. The modified weight matrix U h obtained from (iii) step above, is a solution of the following optimization problem: min U h || ˆ U h -U h || 2 F , s.t. σ 1 ( U h ) < 2 -δ .

7 Evaluation and Experiments

The first two experiments are single pendulum and double pendulum, where the dynamics of the latter one is more complicated. The third experiment is ball bouncing in irregular polygon to demonstrate the ability of the GIN when it faces irregular environments with multiple underlying dynamics. The fourth experiment covers visual odometry task for real world data. The last experiment shows the effectiveness of theorems 3 and 4 for gradient explosion handling. Intuitive python code and training algorithm are in appendix A.11. Detailed explanations regarding hyperparameter optimization, network structure, and empirical solutions aimed at avoiding poor local minima are in appendix A.7.

Single Pendulum and Double Pendulum.

In the pendulum experiment, observations are the image sequences comprising 100 time steps, each has the size of 24 × 24 distorted with time correlated noise. We perform the filtering-smoothing by the GIN. In state estimation task, the log-likelihood and

squared error (SE) of the estimated states are given in table 1. With n = 3 m ( n and m representing state and transferred observation dimensions) as shown in Table 1, intuitively states contain information about position, velocity, and acceleration, increasing the likelihood compared to n = 2 m , where only position and velocity are allocated to states. For the imputation task, we randomly remove half of the images from the generated sequences and perform image imputation by predicting the missing parts. The numerical results are in table 2 and imputed images are in figures 5 and 6. In table 2, the models with informed boolean masks are aware of available and missing images, while uninformed masks use a black image as input for missing images. This encourages the model to accurately infer the dynamics for image generation.

For ablation study, first we compare GIN with simple encoder-decoder without latent parameterization, and

then with LGSSM, where the GRU cells are omitted from the GIN structure and classic filteringsmoothing equations are used instead. We also check the usefulness of filtering-smoothing parameterization by replacing the GIN's transition block with GRU and LSTM cells and directly parameterize output distribution (like AE-RNN [36] and [37]). We compare various VI models, including VAE, IWVAE, VAE-RNN [38] which utilizes VAEs with recurrent cells, and EM-based VI approaches such as SVAE, KVAE, EKVAE and MVAE. Our comparison also encompasses DeepAR [39], an AR model. Our thorough comparison, covering the CRU, RKN, a group of VAEs, and AR models, enhances the depth of our analysis. The GIN consistently outperforms all other models. Additionally, we provide results using the MSE metric to highlight the competitive prediction accuracy of our approach (See appendix A.10).

Bouncing Ball In Irregular Polygon. To demonstrate the adaptability of the GIN, we subjected it to a more complex environment. We generated 1000 sequences, each comprising 70 time steps, depicting a ball moving within an irregular polygon featuring random shapes.

The initial position and velocity of the ball were randomized. No external forces acted upon the ball, and collisions with the walls were elastic. In addition to the aforementioned models, we include latent SLDS structure from the SVAE approach for the comparison. SLDS is particularly relevant in scenarios such as this one, where separate dynamics for each state are considered, as seen when the ball bounces off the wall. We exclude simple baselines, such as encoder-decoder, VAE, and IW-VAE, from the comparison due to their lack of competitive performance. The numerical results are in tables 1 and 2. In Figure 7, we illustrate samples obtained from the trained smoothing distributions of GIN, LGSSM, and KVAE, showcasing their proximity to the ground truth. For this demonstration, we specifically select the 35th time step for sampling, denoted as L ( s 35 | o 1:70 ) . Here, s 1 and s 2 represent the estimated xy position of the ball within the polygon. Generating 20 frames with GIN in four different irregular environments is shown in figure

s1

s2

s1

s2

s1

s2

s1

s2

15

20

25

30

35

s1

p

(

gt

135) =

(

s

1

25)

35|70 =24.91

35|70 =22.83

35|70 =29.15

0

5

10

15

20

p

(

gt

145) =

(

s

1

12)

45|70 =11.86

45|70 =8.14

45|70 =8.09

GIN

LGSSM

KVAE

0

5

10

15

20

s2

p

(

gt

235) =

(

s

2

10)

35|70 =10.07

35|70 =12.51

35|70 =8.64

GIN

LGSSM

KVAE

0

5

10

15

20

p

(

gt

245) =

(

s

2

11)

45|70 =10.73

45|70 =14.11

45|70 =15.25

1. We have created animated files that demonstrate both sequence generation and imputation tasks within irregular polygons. Please visit: https://sites.google.com/view/ginneurips2024.

Visual Odometry of KITTI Dataset. We also evaluate the GIN with the higher dimensional observations for the visual odometry task on the KITTI dataset [40]. This dataset consists of 11 separated image sequences with their corresponding labels. We use a feature extractor network proposed by Zhou et al. in [41] to obtain the transferred observations of the GIN, i.e. ( w , r ). Additionally, we compare the results with AE-RNN(LSTM), AE-RNN(GRU), DeepVO [42] and KVAE. The numerical results are in table 3, where the common evaluation scheme for the KITTI dataset is exploited. The results of the KV AE degrades substantially as we have to reduce the size of the transferred observation to prevent the complexity of matrix inversion. Sampling visualization from smoothing state distribution are in appendix A.9.

Gradient Explosion Experiment. Table 4 lists the log likelihood and its standard deviation for three experiments: single pendulum, double pendulum, and irregular polygon. These experiments were trained under different settings for gradient explosion handling: the conventional Gradient Clipping (GC), our first solution using Singular Value Decomposition (SVD), and our second solution using the Gershgorin Circle Theorem (GCT). The table empirically demonstrates that our method outperforms conventional gradient clipping. In this table, θ represents the threshold for gradient clipping as introduced in [10], and δ is the threshold in our method that keeps the spectral radius less than 2, i.e., σ 1 ( U h ) + δ = 2 . As shown in table 4, gradient clipping failed to train for high θ , while our approach achieves lower perplexity and higher log likelihood, ensuring stability compared to gradient clipping by constraining the GRU to be stable. We provide a comprehensive discussion on the δ variable, including how to tune it in different environments and its impact on model performance, in the appendix A.8.1.

8 Conclusion

This paper introduces the GIN for representation learning in high-dimensional SSMs. The data flow is managed by Bayesian filtering-smoothing, and the use of GRU-based KG and SG networks addresses computational issues, resulting in an efficient and numerically stable model. The GIN learns system dynamics end-to-end, making it a high-performance model with strong system identification capabilities. It also provides insightful uncertainty representations for predictions and outperforms various counterparts, including generative models with variational inference, autoregressive models, and other baselines.

9 Acknowledgments

We thank Yasin Abbasi-Yadkori for useful discussions and suggestions that contributed to this work.

Bibliography

C. M. Bishop and N. M. Nasrabadi, Pattern recognition and machine learning . Springer, 2006, vol. 4, no. 4.

R. M. Neal et al. , 'Mcmc using hamiltonian dynamics,' Handbook of markov chain monte carlo , vol. 2, no. 11, p. 2, 2011.

D. P. Kingma and M. Welling, 'Auto-encoding variational bayes,' arXiv preprint arXiv:1312.6114 , 2013.

D. Koller and N. Friedman, Probabilistic graphical models: principles and techniques . MIT press, 2009.

R. E. Kalman, 'A new approach to linear filtering and prediction problems,' 1960.

H. E. Rauch, F. Tung, and C. T. Striebel, 'Maximum likelihood estimates of linear dynamic systems,' AIAA journal , vol. 3, no. 8, pp. 1445-1450, 1965.

E. A. Wan and R. Van Der Merwe, 'The unscented kalman filter for nonlinear estimation,' in Proceedings of the IEEE 2000 Adaptive Systems for Signal Processing, Communications, and Control Symposium (Cat. No. 00EX373) . Ieee, 2000, pp. 153-158.

L. Ljung, 'Asymptotic behavior of the extended kalman filter as a parameter estimator for linear systems,' IEEE Transactions on Automatic Control , vol. 24, no. 1, pp. 36-50, 1979.

S. Linderman, M. Johnson, A. Miller, R. Adams, D. Blei, and L. Paninski, 'Bayesian learning and inference in recurrent switching linear dynamical systems,' in Artificial Intelligence and Statistics . PMLR, 2017, pp. 914-922.

R. Pascanu, T. Mikolov, and Y. Bengio, 'On the difficulty of training recurrent neural networks,' in International conference on machine learning . Pmlr, 2013, pp. 1310-1318.

K. Doya et al. , 'Bifurcations in the learning of recurrent neural networks 3,' learning (RTRL) , vol. 3, p. 17, 1992.

T. Haarnoja, A. Ajay, S. Levine, and P. Abbeel, 'Backprop kf: Learning discriminative deterministic state estimators,' in Advances in neural information processing systems , 2016, pp. 4376-4384.

R. Krishnan, U. Shalit, and D. Sontag, 'Structured inference networks for nonlinear state space models,' in Proceedings of the AAAI Conference on Artificial Intelligence , vol. 31, no. 1, 2017.

P. G. Chang, K. P. Murphy, and M. Jones, 'On diagonal approximations to the extended kalman filter for online training of bayesian neural networks,' in Continual Lifelong Learning Workshop at ACML 2022 , 2022.

M. Jones, T. R. Scott, M. Ren, G. F. Elsayed, K. Hermann, D. Mayo, and M. C. Mozer, 'Learning in temporally structured environments,' in The Eleventh International Conference on Learning Representations , 2022.

M. Watter, J. T. Springenberg, J. Boedecker, and M. Riedmiller, 'Embed to control: A locally linear latent dynamics model for control from raw images,' arXiv preprint arXiv:1506.07365 , 2015.

Y. Burda, R. Grosse, and R. Salakhutdinov, 'Importance weighted autoencoders,' arXiv preprint arXiv:1509.00519 , 2015.

M. J. Johnson, D. K. Duvenaud, A. Wiltschko, R. P. Adams, and S. R. Datta, 'Composing graphical models with neural networks for structured representations and fast inference,' Advances in neural information processing systems , vol. 29, 2016.

M. Fraccaro, S. Kamronn, U. Paquet, and O. Winther, 'A disentangled recognition and nonlinear dynamics model for unsupervised learning,' arXiv preprint arXiv:1710.05741 , 2017.

Y. Li and S. Mandt, 'Disentangled sequential autoencoder,' arXiv preprint arXiv:1803.02991 , 2018.

A. Klushyn, R. Kurle, M. Soelch, B. Cseke, and P. van der Smagt, 'Latent matters: Learning deep state-space models,' Advances in Neural Information Processing Systems , vol. 34, 2021.

F. Tonolini, N. Aletras, Y. Jiao, and G. Kazai, 'Robust weak supervision with variational auto-encoders,' in International Conference on Machine Learning . PMLR, 2023, pp. 34 39434 408.

H. Zhu, C. Balsells-Rodas, and Y. Li, 'Markovian gaussian process variational autoencoders,' in International Conference on Machine Learning . PMLR, 2023, pp. 42 938-42 961.

P. Becker, H. Pandya, G. Gebhardt, C. Zhao, C. J. Taylor, and G. Neumann, 'Recurrent kalman networks: Factorized inference in high-dimensional deep feature spaces,' in International Conference on Machine Learning . PMLR, 2019, pp. 544-552.

M. Schirmer, M. Eltayeb, S. Lessmann, and M. Rudolph, 'Modeling irregular time series with continuous recurrent units,' in International Conference on Machine Learning . PMLR, 2022, pp. 19 388-19 405.

S. Hochreiter and J. Schmidhuber, 'Long short-term memory,' Neural computation , vol. 9, no. 8, pp. 1735-1780, 1997.

K. Cho, B. Van Merriënboer, D. Bahdanau, and Y. Bengio, 'On the properties of neural machine translation: Encoder-decoder approaches,' arXiv preprint arXiv:1409.1259 , 2014.

R. Wilson and L. Finkel, 'A neural implementation of the kalman filter,' Advances in neural information processing systems , vol. 22, pp. 2062-2070, 2009.

I. Buchnik, G. Revach, D. Steger, R. J. Van Sloun, T. Routtenberg, and N. Shlezinger, 'Latentkalmannet: Learned kalman filtering for tracking from high-dimensional signals,' IEEE Transactions on Signal Processing , 2023.

G. Revach, N. Shlezinger, X. Ni, A. L. Escoriza, R. J. van Sloun, and Y. C. Eldar, 'Kalmannet: Neural network aided kalman filtering for partially known dynamics,' arXiv preprint arXiv:2107.10043 , 2021.

D. Ruhe and P. Forré, 'Self-supervised inference in state-space models,' arXiv preprint arXiv:2107.13349 , 2021.

T. Liu, K. A. Lee, Q. Wang, and H. Li, 'Disentangling voice and content with self-supervision for speaker recognition,' Advances in Neural Information Processing Systems , vol. 36, pp. 50 221-50 236, 2023.

M. P. Deisenroth and H. Ohlsson, 'A general perspective on gaussian filtering and smoothing: Explaining current and deriving new algorithms,' in Proceedings of the 2011 American Control Conference . IEEE, 2011, pp. 1807-1812.

P. Gilabert, G. Montoro, and E. Bertran, 'On the wiener and hammerstein models for power amplifier predistortion,' in 2005 Asia-Pacific Microwave Conference Proceedings , vol. 2. IEEE, 2005, pp. 4-pp.

J. Chung, C. Gulcehre, K. Cho, and Y. Bengio, 'Empirical evaluation of gated recurrent neural networks on sequence modeling,' arXiv preprint arXiv:1412.3555 , 2014.

N. Srivastava, E. Mansimov, and R. Salakhudinov, 'Unsupervised learning of video representations using lstms,' in International conference on machine learning . PMLR, 2015, pp. 843-852.

H. Li, S. Yu, and J. Principe, 'Causal recurrent variational autoencoder for medical time series generation,' in Proceedings of the AAAI Conference on Artificial Intelligence , vol. 37, no. 7, 2023, pp. 8562-8570.

J. Chung, K. Kastner, L. Dinh, K. Goel, A. C. Courville, and Y. Bengio, 'A recurrent latent variable model for sequential data,' Advances in neural information processing systems , vol. 28, 2015.

D. Salinas, V. Flunkert, J. Gasthaus, and T. Januschowski, 'Deepar: Probabilistic forecasting with autoregressive recurrent networks,' International Journal of Forecasting , vol. 36, no. 3, pp. 1181-1191, 2020.

A. Geiger, P. Lenz, and R. Urtasun, 'Are we ready for autonomous driving? the kitti vision benchmark suite,' in 2012 IEEE conference on computer vision and pattern recognition . IEEE, 2012, pp. 3354-3361.

T. Zhou, M. Brown, N. Snavely, and D. G. Lowe, 'Unsupervised learning of depth and egomotion from video,' in Proceedings of the IEEE conference on computer vision and pattern recognition , 2017, pp. 1851-1858.

S. Wang, R. Clark, H. Wen, and N. Trigoni, 'Deepvo: Towards end-to-end visual odometry with deep recurrent convolutional neural networks,' in 2017 IEEE international conference on robotics and automation (ICRA) . IEEE, 2017, pp. 2043-2050.

K. Sohn, H. Lee, and X. Yan, 'Learning structured output representation using deep conditional generative models,' Advances in neural information processing systems , vol. 28, 2015.

D. Angeli, E. D. Sontag, and Y. Wang, 'A characterization of integral input-to-state stability,' IEEE Transactions on Automatic Control , vol. 45, no. 6, pp. 1082-1097, 2000.

R. A. Horn and C. R. Johnson, Matrix analysis . Cambridge university press, 2012.

D. P. Kingma and J. Ba, 'Adam: A method for stochastic optimization,' arXiv preprint arXiv:1412.6980 , 2014.

P. J. Werbos, 'Backpropagation through time: what it does and how to do it,' Proceedings of the IEEE , vol. 78, no. 10, pp. 1550-1560, 1990.

J. L. Ba, J. R. Kiros, and G. E. Hinton, 'Layer normalization,' arXiv preprint arXiv:1607.06450 , 2016.

J. M. Wang, D. J. Fleet, and A. Hertzmann, 'Gaussian process dynamical models for human motion,' IEEE transactions on pattern analysis and machine intelligence , vol. 30, no. 2, pp. 283-298, 2007.

J. Ko and D. Fox, 'Learning gp-bayesfilters via gaussian process latent variable models,' Autonomous Robots , vol. 30, no. 1, pp. 3-23, 2011.

R. Frigola, F. Lindsten, T. B. Schön, and C. E. Rasmussen, 'Bayesian inference and learning in gaussian process state-space models with particle mcmc,' Advances in neural information processing systems , vol. 26, 2013.

M. Schoukens and K. Tiels, 'Identification of block-oriented nonlinear systems starting from linear approximations: A survey,' Automatica , vol. 85, pp. 272-292, 2017.

R. Li, S. John, and A. Solin, 'Improving hyperparameter learning under approximate inference in gaussian process models,' in International Conference on Machine Learning . PMLR, 2023, pp. 19 595-19 615.

N. Wahlström, T. B. Schön, and M. P. Deisenroth, 'From pixels to torques: Policy learning with deep dynamical models,' arXiv preprint arXiv:1502.02251 , 2015.

E. Archer, I. M. Park, L. Buesing, J. Cunningham, and L. Paninski, 'Black box variational inference for state space models,' arXiv preprint arXiv:1511.07367 , 2015.

M. Karl, M. Soelch, J. Bayer, and P. Van der Smagt, 'Deep variational bayes filters: Unsupervised learning of state space models from raw data,' arXiv preprint arXiv:1605.06432 , 2016.

C. Naesseth, S. Linderman, R. Ranganath, and D. Blei, 'Variational sequential monte carlo,' in International conference on artificial intelligence and statistics . PMLR, 2018, pp. 968-977.

S. S. Rangapuram, M. W. Seeger, J. Gasthaus, L. Stella, Y. Wang, and T. Januschowski, 'Deep state space models for time series forecasting,' Advances in neural information processing systems , vol. 31, 2018.

V. Garcia Satorras, Z. Akata, and M. Welling, 'Combining generative and discriminative models for hybrid inference,' Advances in Neural Information Processing Systems , vol. 32, 2019.

A Appendix

A.1 Filtering and Smoothing Parameterization

The Kalman filter operates through two iterative steps: prediction and update. In the prediction step, the filter uses prior state information to make an estimate, while in the update step, it adjusts this estimate based on new observations. Assuming normally distributed additive process and observation noise, the filter can effectively perform these steps. During the prediction step, the filter employs the transition matrix F t +1 to estimate the next priors p γ 1: t +1 ( x t +1 | w 1: t ) = ( µ t +1 | t , Σ t +1 | t ) , which represent the next state estimates without incorporating any new observations.

When a new observation is available, the Kalman filter proceeds to the second step, where it updates the predicted prior using the new observation and the emission matrix H t +1 . This results in the next posterior, ( µ t +1 | t +1 , Σ t +1 | t +1 ) .

The entire observation update procedure can be viewed as a weighted mean between the next prior, derived from the state update, and the new observation. This weighting is influenced by Q and R , reflecting their inherent uncertainties. We derive the smoothing parameterization by leveraging the Markov property, which states that x t is independent of future observations w t +1: T given x t +1 . Although x t +1 is unknown, there is a distribution over it. By conditioning on x t +1 and then marginalizing out, we can obtain x t conditioned on w 1: T .

By using induction and and smoothed distribution for t +1 :

we calculate the filtered two-slice distribution as follows:

by using Gaussian conditioning we have:

where J t = Σ t | t F t +1 [ Σ t +1 | t ] -1 . We calculate the smoothed distribution for t using the rules of iterated expectation and covariance:

A.2 Process Noise Matrix

As stated in (13), we can elaborate the process noise matrix at time t in more details

combining (13) into (22) results in

which is a function of F t , Q t -1 , F t -1 and H t -1 . In the GIN, ˆ F t and ˆ H t are learned by the Dynamic Network with the input of x t -1 . From (15), µ t -1 | t -1 is derived as a function of both F t -1 and H t -1 , meaning the learned ˆ F t carries the information of both H t -1 and F t -1 . Therefore, one can rewrite the equation (23) as

where g is a nonlinear function mapping µ t -1 | t -1 and Q t -1 to Q t and the graphical model for such choice of structure is in appendix figure 8b. It is possible to go one step further and simplify µ t -1 | t -1 more, as it has Σ t -1 | t -1 term in (15), combining it with (13) results in:

indicating that not only F t -1 and H t -1 , but also Q t -1 is included in µ t -1 | t -1 , meaning that Q t can be written solely as a function of µ t -1 | t -1 and the graphical model for such choice is in figure 8a.

We refer to g as the Q Network , which can be modeled either by an MLP, as shown in (26), or by a recurrent network, as detailed in (24).

A.3 Output Distribution

In the case of grayscale images, each pixel y is either one or zero with the probability p or 1 -p , respectively, meaning that P ( Y = y ) = p y (1 -p ) 1 -y . This probability equation can be rewritten in the form of an exponential family as follows:

𝑞𝑡-1

x𝑡-1

w𝑡-1

𝑞𝑡

x𝑡

w𝑡

𝑞𝑡+1

x𝑡+1

w𝑡+1

…

…

…

…

…

…

o𝑡-1

s 𝑡-1

o𝑡

s

𝑡

o𝑡+1

s 𝑡+1

𝑞𝑡-1

x𝑡-1

w𝑡-1

𝑞𝑡

x𝑡

w𝑡

𝑞𝑡+1

x𝑡+1

w𝑡+1

…

…

…

…

…

…

o𝑡-1

s 𝑡-1

o𝑡

s

𝑡

o𝑡+1

s 𝑡+1

By choosing θ = log ( p 1 -p ) and ψ ( θ ) = log (1 -p ) and h ( y ) = 1 , we can obtain p = 1 1+ e -θ . This means that by considering θ as the last layer of the decoder and applying a sigmoid layer, p is obtained and one can optimize p such that the likelihood is maximized as we did in in (11).

Similarly, consider x as the ground truth state, ˆ x θ as the estimated state, and θ as the model variables. The residual follows a Gaussian distribution: x = ˆ x θ + ϵ ∼ N (ˆ x θ , ˆ σ θ ) , where ˆ σ θ is the estimated variance. Then, the negative log likelihood is given by (28) as we obtained it in(10).

A.4 Proof of theorems

We restate the theorems for the ease of readability and then provide the proofs.

A.4.1 Proof of theorem 1

Theorem. The log likelihood of the states given the original observations is determined as:

where the d m ( . ) and d c ( . ) determines those parts of the decoder that obtain the state mean and variance, respectively. N sequences of ( x ( i ) 1: T , w ( i ) 1: T ) ∼ q γ 1: T ( x 1: T , w 1: T | o 1: T ) are sampled for Monte Carlo integration estimation.

Proof.

The inequality in (32) arises from the positivity of the Kullback-Leibler (KL) term. The approximation in (34) results from assuming equality between the prior and the approximated posterior, i.e., p γ 1: T ( w 1: T , x 1: T | o 1: T ) = q ϕ ( x 1: T , w 1: T | o 1: T , s 1: T ) . This assumption is not very restrictive as explained further in [43]. The approximation in (35) is coming from the equality assumption of q ϕ ( x 1: T , w 1: T | o 1: T , s 1: T ) and q γ 1: T ( x 1: T , w 1: T | o 1: T ) . This assumption implies that given the original observation o 1: T , our latent representation ( x 1: T , w 1: T ) is independent of the ground truth states s 1: T . In simpler terms, s 1: T is concealed within o 1: T , allowing us to disregard s 1: T . The approximation in (37) is because of Monte Carlo integration estimation.

A.4.2 Proof of theorem 2

Theorem. The (lower bound of) log likelihood of the original images is:

with N sequences of ( x ( i ) 1: T , w ( i ) 1: T ) ∼ q γ 1: T ( x 1: T , w 1: T | o 1: T ) for Monte Carlo integration estimation. d k ( . ) defines the corresponding part of the decoder that maps the k -th pixel of o t image.

Proof.

The inequality in (44) arises from the positivity of the Kullback-Leibler (KL) term. The approximation in (46) results from assuming equality between the prior and the approximated posterior, i.e., p γ 1: T ( w 1: T , x 1: T ) = q ϕ ( x 1: T , w 1: T | o 1: T ) . Although it is feasible to optimize (44) directly, we found out that by setting prior p γ 1: T ( x 1: T , w 1: T ) equal to approximated posterior q γ 1: T ( x 1: T , w 1: T | o 1: T ) , more stable results with concise improvement are achieved from the numerical results (similar to [19, 21, 43]). The approximation in (48) is because of Monte Carlo integration estimation.

A.4.3 Proof of theorem 3

First we provide a brief review on the GRU cell equations from [35]. Then we introduce two lemmas, by which we can proof theorem 2.

GRU Cell Review. The GRU cell holds the following equations:

where x t and h t are the input and the hidden state vectors, respectively. U xz , U z , U xr , U r , U xh and U h are all weight matrices.

Lemma 1. A GRU cell has a fixed point at h ∗ = 0 , if the input variable x t = 0 .

Proof. After substituting x t = 0 and h t -1 = 0 into (51) and (52), the update gate z t and reset gate r t are both 1 2 . Substituting these values along with x = 0 and h t -1 = 0 into (54) gives the candidate

state ˆ h t = 0 . Finally, substituting h t -1 = 0 , z t = 1 2 , and ˆ h t = 0 into (53) results in the new state h t = 0 . Thus, h t = h t -1 = 0 , indicating that GRU has a fixed point h ∗ = 0 .

Lemma 2. Let I be an h × h identity matrix, λ i ( . ) shows the i -th largest absolute eigenvalue, and A = 1 4 U h + 1 2 I . When the spectral radius | λ 1 ( A ) | < 1 , a GRU cell with input x t = 0 . can be approximated by the following linearized GRU near h t = 0 :

and the fixed point h ∗ = 0 is locally stable.

Proof. By employing Taylor series expansion and linearizing h t with respect to h t -1 at h t -1 = x t = 0 , A in (55) is obtain as 1 4 U h + 1 2 I . Then, from (55) we get that h t = A t h 0 . Since A t is dependent on the t -th powers of the eigenvalues of A , the behavior of the linearized GRU is determined by the eigenvalues of A . According to the Hartman-Grobman theorem, the behavior of a dynamical system near a hyperbolic fixed point is homeomorphic to the behavior of the linearized system. Therefore, when | λ 1 ( A ) | < 1 , the fixed point h ∗ = 0 becomes a hyperbolic fixed point. Thus, a GRU cell can be approximated as h t = A t h 0 . Then, local stability is determined by the spectral radius | λ 1 ( A ) | < 1 and then we have | h t -h 0 | t →∞ = 0 and the fixed point h ∗ = 0 is locally stable.

Theorem 3. When σ 1 ( U h ) < 2 , a GRU cell is locally stable at a fix point h ∗ = 0 .

Proof. Relying on Lemma1 and Lemma2 , we need to satisfy | λ 1 ( 1 4 U h + 1 2 I ) | < 1 . | λ 1 ( 1 4 U h + 1 2 I ) | = | 1 4 λ 1 ( U h ) + 1 2 | due to the properties of eigenvalues. From the triangle inequality we also have | 1 4 λ 1 ( U h ) + 1 2 | ≤ 1 4 | λ 1 ( U h ) | + 1 2 . Using Weyl's inequality for its singular and eigenvalues, we have | λ 1 ( U h ) | ≤ σ 1 ( U h ) and therefore 1 4 | λ 1 ( U h ) | + 1 2 ≤ 1 4 | σ 1 ( U h ) | + 1 2 . Thus, we need to satisfy 1 4 | σ 1 ( U h ) | + 1 2 < 1 condition and conclude 1 4 | λ 1 ( U h ) | + 1 2 ≤ 1 4 | σ 1 ( U h ) | + 1 2 < 1 . This can be achieved by σ 1 ( U h ) < 2 .

The GRU cell is considered as a nonlinear dynamical system with ˙ h t = f ( h t , x t ) model, where one can drive f from (51)-(54). We have shown the conditions for local stability of ˙ h = f ( h , 0 ) in the last paragraph. Now we show ˙ h = f ( h , x ) is locally stable if x t → 0 as t →∞ .

As we have the stability condition for ˙ h = f ( h , 0 ) , there exists V ( h ) that is a Lyapunov function for this system such that:

for a positive definite W ( h ) . When the input variable x is not zero, the system equation becomes ˙ h = f ( h , x ) and the derivative of V ( h ) becomes

and ϵ → 0 as x t → 0 , utilizing Taylor series expansion of f ( h , x ) near x = 0 [44]. Then, as x t → 0 , for any ϵ ≥ 0 , there exists a T such that for all t ≥ T , | x t | < ϵ and therefore we have:

for sufficiently small ϵ , ˙ V ( h ) remains negative.

Therefore, we need to demonstrate that the inputs of the GRU cells in (6) and (8) approach 0 as t → ∞ . As discussed in the main script, we apply a Wishart prior distribution on the covariance matrices, guiding them towards a scaled identity matrix. Over time, this scale diminishes towards zero, causing the covariance matrices-and consequently, the zero bias Conv( . ) operators in (6) and but for a small x t , we can write:

(8)-to converge to 0 as t → ∞ . This implies that the system tends to behave deterministically after prolonged observation, while the stability of the GRU cell is guaranteed. Informally, Theorem 3 requires us to asymptotically shrink the input covariance matrices of our GRU cells toward zero. While this shrinkage is reasonable as we expect the system to behave deterministically in the long run, it imposes a limitation on posterior inference. We acknowledge this as an open problem for future research.

A.4.4 Proof of theorem 4

Theorem. The modified weight matrix U h obtained from (iii) step above, is a solution of the following optimization problem: min U h || ˆ U h -U h || 2 F , s.t. σ 1 ( U h ) < 2 -δ .

Proof. Relying on [45], due to the matrices properties we have:

Assuming that the first s singular values of ˆ U h are greater than 2, we have ∑ h i =1 [ σ i ( ˆ U h ) -σ i ( U h ) ] = ∑ s i =1 [ σ i ( ˆ U h ) -(2 -δ ) ] . According to (61), we need to show that ∑ s i =1 [ σ i ( ˆ U h ) -(2 -δ ) ] 2 is equal to its upper bound || ˆ U h -U h || 2 F , i.e. the optimum solution.

Using the singular value decomposition properties, we have || ˆ U h -U h || 2 F = || WΛV -WΛV || 2 F = || W ( Λ -Λ ) V || 2 F = tr ( V ∗ ( Λ -Λ ) W ∗ W ( Λ -Λ ) V ) = tr ( V ∗ ( Λ -Λ )( Λ -Λ ) V ) = tr ( VV ∗ ( Λ -Λ )( Λ -Λ )) = tr (( Λ -Λ )( Λ -Λ )) = ∑ s i =1 [ σ i ( ˆ U h ) -(2 -δ ) ] 2 .

A.4.5 Proof of theorem 5

Theorem. We can satisfy the constraint introduced in Theorem 3 by applying the Gershgorin circle theorem as follows.

̸

(i) In the i -th row of ˆ U h , sum up the non diagonal elements as i d = ∑ h j =1 ,j = i | ˆ U h ( i, j ) | . Do (i) for the all rows.

(ii) For each ˆ U h ( i, i ) assign a circle cluster centered at ˆ U h ( i, i ) with i d radius in the complex plane.

(iii) Merge the clusters which have intersection and repeat step (iii) until no cluster can be merged any more. Finally the set of clusters C = {C 1 , ..., C k } with length of k ≤ h is obtained.

(iv) Define the radius of C j cluster as d C j = max ( ˆ U h ( i, i )+ i d |∀ ˆ U h ( i, i ) ∈ C j ) . Calculate the radius of all k clusters in step (iv).

(v) For each cluster C j , if ( ˆ U h ( i, i ) + d C j |∀ ˆ U h ( i, i ) ∈ C j ) > 2 , shift ˆ U h ( i, i ) ← (2 -d C j -δ ) for a small δ . Otherwise, keep ˆ U h ( i, i ) unchanged. Do step (v) for all clusters and their elements, i.e. ˆ U h ( i, i ) .

These five steps, ensures the new ˆ U h follows the constraint introduced in Theorem 3 .

Proof. The five steps elaborated above are direct results of the generalized Gershgorin circle theorem. This approach is faster than SVD method introduced in the paper, while it is not the best solution of the optimization problem introduced in the Theorem 4 .

A.5 Discussion About KL Term of Mode Collapse

The term introduced in the mode collapse handling section is:

Here p i ( ˇ F ) ∼ MN n × n ( ˇ F i , I , I ) and p j ( ˇ F ) ∼ MN n × n ( ˇ F j , I , I ) , are matrix normal distributions with priors pri i ( ˇ F ) ∼ MN n × n ( Diag ( m i ) , I , I ) and pri j ( ˇ F ) ∼ MN n × n ( Diag ( m j ) , I , I ) , respectively. We can expand (62) as:

Relying on [45], we can simplify each term in (65) as:

Thus, we can rewrite (62) as:

which is half of (12) plus a constant.

A.6 Noise Generation Process

In our experiments, we use a time-correlated noise generation scheme to demonstrate the system's noise robustness. This method introduces a sequence of factors, f t , of the same length as the data sequence, making the noise factors correlated over time. Let f 0 ∼ U (0 , 1) and f t +1 = min ( max (0 , f t + r t ) , 1) with r t ∼ U ( -0 . 2 , 0 . 2) , where f 0 is the initialized factor and U is the uniform distribution. Two thresholds, t 1 ∼ U (0 , 0 . 25) and t 2 ∼ U (0 . 75 , 1) , are defined. Values of f t less than t 1 are set to 0, values greater than t 2 are set to 1, and the remaining values are linearly scaled within the range [0 , 1] . The original observation at time t , o t , is then given by o t = f t i t +(1 -f t ) i pn t , where i t is the true image at time t and i pn t is the pure noise generated for time t .

A.7 Hyperparameters and training

In all experiments, we used the Adam optimizer [46] on an NVIDIA GeForce GTX 1050 Ti with 16GB RAM. To ensure optimal hyperparameters, we conducted a grid search. We searched for the initial learning rate in the range of 0 . 001 to 0 . 2 with increments of 0 . 005 , selecting the one that yielded the highest log-likelihood. We chose an initial learning rate of 0 . 006 with an exponential decay rate of 0 . 9 every 10 epochs. Backpropagation through time [47] was employed to compute gradients, given the use of GRU cells in our structure. The gradients are applied to GRU cells with the constraint explained in theorem 3, where we use the spectral theorem mentioned in the main script . We applied the layer normalization technique [48] to stabilize the dynamics within the recurrent structure and normalize the filter response. The Elu + 1 activation function was used to ensure the positivity of the diagonal elements of the process noise and covariance matrices.

Dynamics

Network

𝐱 0:𝑇-1|0:𝑇-1

Prediction

( ෠

𝐅 1:𝑇

, ෡

𝐇1:𝑇

)

(ෝ

𝝁 1:𝑇|0:𝑇-1

, ෡

𝚺 1:𝑇|0:𝑇-1

)

𝐺𝑅𝑈 𝐾𝐺

Filtering Step

𝒓1:𝑇

𝐌1:𝑇

𝒘1:𝑇

(ෝ

𝝁 1:𝑇|1:𝑇

, ෡

𝚺 1:𝑇|1:𝑇

)

Smoothing Step

𝐺𝑅𝑈 𝑆𝐺

(ෝ

𝝁 0:𝑇-1|0:𝑇-1

, ෡

𝚺 0:𝑇-1|0:𝑇-1

)

𝐍1:𝑇

෡

𝐇1:𝑇

෠

𝐅 1:𝑇

(ෝ

𝝁 1:𝑇|𝑇

, ෡

𝚺 1:𝑇|𝑇

)

Encoder

𝒐1:𝑇

Prediction

Filtering

(𝒓 1:𝑇

, 𝒘

1:𝑇

)

(ෝ

𝝁1:𝑇|1:𝑇

, ෡

𝚺1:𝑇|1:𝑇

)

Decoder

{𝒐 1:𝑇

, 𝒔

1:𝑇

}

Smoothing

(ෝ

𝝁1:𝑇|𝑇

, ෡

𝚺1:𝑇|𝑇

)

GIN Cell

Conv

2D

෡

𝚺 1:𝑇|0:𝑇-1

Conv

2D

(ෝ

𝝁 1:𝑇|0:𝑇-1

, ෡

𝚺 1:𝑇|0:𝑇-1

)

෡

𝚺 1:𝑇|0:𝑇-1

(ෝ

𝝁 1:𝑇|1:𝑇

, ෡

𝚺 1:𝑇|1:𝑇

)

(ෝ

𝝁1:𝑇|0:𝑇-1

, ෡

𝚺1:𝑇|0:𝑇-1

)

To prevent the model from getting stuck in poor local minima, such as focusing excessively on reconstruction rather than learning the dynamics obtained through filtering-smoothing, we find it beneficial to employ two training techniques for end-to-end learning:

1Generating time correlated noisy sequences as consecutive observations, forces the model to learn the dynamics instead of focusing on reconstruction, e.g. figures 11 and 13 in appendix.

2During the initial epochs, the system focuses on learning the auto-encoder and globally learned parameters, such as ˇ F k and ˇ H k , while excluding the Dynamic Net parameters α t ( x t -1 ) . Afterwards, all parameters are jointly learned. This approach enables the system to initially acquire good embeddings and meaningful latent vectors before integrating the learning of K different dynamics variables.

In our experiments, we set K = 15 to accommodate the complexity of the latent space, with each latent representing a unique dynamical scenario. Generally, parameter tuning is not challenging when the GIN is sufficiently flexible, as it can learn to prune unused elements through the Dynamic Net .

A.8 Proposed architecture, qualitative comparison with the SOTA and empirical complexity analysis

A.8.1 Architectures

The proposed detailed structure is shown in figure 9. To design the dynamic network, we use an MLP with 60 hidden units and a ReLU activation function, and a softmax activation in the last layer. The state mean, with a size of n , is the input to the dynamic network, which outputs k coefficients. The structures of the encoder and decoder are shown in Table 5. In this table, m represents the transferred observation dimension, with various values considered in the results. For state estimation tasks, the output size ( out ) is 4, 4, 8 and 12 for the single-pendulum, bouncing ball, double-pendulum and KITTI visual odometry experiments, respectively. For the imputation task, the number of hidden units in the KG and SG networks is set to 40 and 30, respectively. The convolutional layer applied over the covariance matrix has 8 filters with a kernel size of 5 and zero bias.

Gradient Explosion Experiment.

Referring to table 4, higher values of δ deteriorate performance because the convergence speed of the linearized GRU hidden state depends on σ 1 ( U h ) (see Eq. (55) in the appendix), which is projected to 2 -δ . Therefore, longer time dependencies in the data are eliminated for high δ . On the other hand, smaller values of δ close to zero make the model more prone to noise [10]. Therefore, it is necessary to tune an appropriate δ within the range 0 < δ < 2 , which is much easier than tuning the unbounded θ .

A.8.2 Empirical running times and parameters

We present the number of parameters for the utilized cell structures in our experiments and their corresponding empirical running times for 1 epoch in Table 6. In the first row of each model structure, we set the number of parameters approximately equal to our GIN to demonstrate the GIN's superior performance with the same parameter count. The extra running time of EM-variational approaches, like KVAE, is due to the use of classic Bayesian equations, which significantly increase running time for higher-dimensional observations. However, the GIN avoids this issue. The number of parameters in the GIN is noticeably lower than in other memory cells, such as LSTM and GRU, and EM-variational methods. This efficiency is achieved by converting high-dimensional sparse covariance matrices into lower-dimensional covariance matrices using a convolutional operator.

A.8.3 Qualitative comparison of the GIN to recent related work.

Switching Linear Dynamical System models (SLDS) decompose complex, nonlinear time series data into sequences of simpler, reusable dynamical modes. Fitting an SLDS to data enables us to learn a flexible nonlinear generative model and parse data sequences into coherent discrete units. [9] is a SLDS model that incorporates additional latent variables to switch among different linear dynamics. However, this approach relies on Gibbs sampling for inferring the parameters, rendering it impractical for large datasets due to scalability issues.

Auto-regressive state space models (ARSSMs) are a class of dynamic models commonly used in time series analysis and forecasting. In these models, the evolution of a system over time is described by a set of states, by observing all data. Auto-Regressive (AR) Hidden Markov Models (AR-HMM) explain time series structures by defining a mapping from past observations to the current observation. [39] presents an ARHMM approach where target values are used directly as inputs. However, this reliance on target values makes the model more susceptible to noise.

Toward learning state space (system identification), a number of works including those by [4953] have proposed algorithms for learning Gaussian Process SSMs (GPSSMs) through maximum likelihood estimation using the iterative EM algorithm. where filtering-smoothing with a set of fixed parameters γ is conducted (E step), and then updating the set of parameters γ is performed such that the obtained likelihood is maximized (M step). Frigola et al. [51] obtain sample trajectories from the smoothing distribution and then, conditioned on this trajectory, conduct the M step for the model's parameters. In the context of System Identification (SI), the GIN is akin to a Hammerstein-Wiener (HW) model [34], as it estimates system parameters directly from observations while transferring the observations through nonlinear functions.

In Table 7, we compare various algorithms to evaluate their ability to handle high-dimensional observations, learn dynamics, estimate state accurately, provide uncertainty estimates while handling noisy data, manage missing data, and perform direct optimization. Classic LGSSMs, such as the EKF and UKF, rely on the linearization of transition and emission equations and apply classic Bayesian updates over the linearized system with respect to the states. In other words, classic LGSSMs are model-based. In contrast, the GIN uses a data-driven network to learn parameters, offering a different approach from the traditional model-based methods.

A.9 Visualization and The Imputation

Graphical results of informed, uninformed and noisy observations for image imputation task for both single and double pendulum experiments can be found in 10, 11, 12 and 13 figures. Inference for the trained smoothing distribution of all experiments are in 14, 15, 16, , 17, 18, 19, 20, 21, 22, 23 and 24 figures, where we generated samples from the smoothing estimated states distribution, ( s t | o 1: T ) . Then we fit density on the generated samples. This visualization shows the effectiveness of the GIN in reducing the uncertainty of the estimates compare to LGSSM and KVAE.

Ground Truth

LGSSM(filter)

LGSSM(smooth)

GIN(filter)

Ground Truth

LGSSM(filter)

LGSSM(smooth)

GIN(filter)

业

GIN(smooth)

GIN(smooth)

Ground Truth

LGSSM(filter)

LGSSM(smooth)

GIN(filter)

Ground Truth

LGSSM(filter)

LGSSM(smooth)

GIN(filter)

GIN(smooth)

GIN(smooth)

GIN

LGSSM

KVAE

0.2

0.4

0.6

0.8

1.0

1.2

s1

p

(

gt

1100) =

(

s

1

0.7)

mean = 0.69

mean = 0.93

mean = 0.53

GIN

LGSSM

KVAE

0.2

0.4

0.6

0.8

1.0

1.2

s2

p

(

gt

2100) =

(

s

2

0.7)

mean = 0.73

mean = 0.91

mean = 0.42

ଵ଴଴

ଵ଴଴

(a) GIN

ଵ଴଴

ଵ଴଴

(b) LGSSM

ଵ଴଴

ଵ଴଴

(c) KVAE

}.0.1

2

S

0.8

0.6

0.4

0.4

0.6

S1

0.8

1.0

1.2

a.8

S

0.6

0.4

0.4

0.6

S1

0.8

1.0

1.2

0.6

0.4

0.4

0.6

S1

0.8

1.0

1.2

GIN

LGSSM

KVAE

0.2

0.0

0.2

0.4

0.6

0.8

s1

p

(

gt

1100) =

(

s

1

0.35)

mean = 0.36

mean = 0.38

mean = 0.44

GIN

LGSSM

KVAE

0.2

0.0

0.2

0.4

0.6

0.8

s2

p

(

gt

2100) =

(

s

2

0.35)

mean = 0.38

mean = 0.28

mean = 0.29

ଵ଴଴

ଵ଴଴

(a) GIN

ଵ଴଴

ଵ଴଴

(b) LGSSM

ଵ଴଴

ଵ଴଴

(c) KVAE

0.7

0.6

0.5

0.4

E'0

0.2

0.1

- 0°0

0.1

0.0

0.1

0.2

S1

0.3

0.4

0.5

0.6

0.7

0.7

0.6 -

0.5

0.4 -

0.3

0.2

0.1 -

0.0

0.1

0.1

0.0

0.1

0.2

s1

0.3

0.4

0.5

0.6

0.7

0.7

0.6

- S'0

0.4

0.3

0.2

0.1

00

0.1

0.1

0.0

0.1

0.2

S1

0.3

0.4

0.5

0.6

0.7

GIN

LGSSM

KVAE

0.6

0.8

1.0

1.2

1.4

s3

p

(

gt

3100) =

(

s

3

1)

mean = 1.0

mean = 0.94

mean = 0.96

GIN

LGSSM

KVAE

0.4

0.2

0.0

0.2

0.4

s4

p

(

gt

4100) =

(

s

4)

mean = -0.01

mean = 0.08

mean = 0.1

ଵ଴଴

ଵ଴଴

(a) GIN

ଵ଴଴

ଵ଴଴

(b) LGSSM

ଵ଴଴

ଵ଴଴

(c) KVAE

0.4

0.2

00

0.2

0.4

0.6

0.8

S3

1.0

1.2

1.4

0.4

0.2

$

00

0.2

0.4

0.6

0.8

S3

1.0

1.2

1.4

0.4

0.2

$

0.0

0.2

0.4

0.6

0.8

S3

1.0

1.2

1.4

60.0

57.5

55.0

52.5

50.0

47.5

45.0

42.5

40.0

s1

p

(

gt

1100) =

(

s

1+50)

100|500 =-50.12

100|500 =-52.25

100|500 =-43.99

GIN

LGSSM

KVAE

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

s2

p

(

gt

1100) =

(

s

1

10)

100|500 =10.08

100|500 =12.86

100|500 =8.46

A.10 MSE Results for The State Estimation and Estimated KG-SG

The MSE results for the single pendulum, double pendulum, and bouncing ball experiments are shown in Tables 8 and 9 and 10. In addition to the dynamics equation (5), where the ˆ F matrix includes the effects of the process noise, two other solutions introduced in Section 5 are included in the MSE results. Parameterizing the process noise as in (24) with a GRU cell is indicated by GRU( Q ), while another parameterization, shown as MLP( Q ), is done using (26).

ଵ଴଴

ଵ଴଴

(a) GIN

ଵ଴଴

ଵ଴଴

(b) LGSSM

ଵ଴଴

ଵ଴଴

(c) KVAE

20.0

17.5

15.0

12.5

2

S

10.0

7.5

5.0

2.5

0.0 +

S1

20.0

17.5

12.5

2

S

10.0

7.5

5.0

2.5

0.0+

S1

20.0

17.5

12.5

S

2

10.0

7.5

5.0+

0.0+

S1

A.11 Algorithms and python inference code

Algorithm Training the GIN

Input: Observations o 1: T , last posteriors (ˆ µ 0: T -1 | 0: T -1 , ˆ Σ 0: T -1 | 0: T -1 ) ( x 0: T -1 | 0: T -1 ) = Sample ( (ˆ µ 0: T -1 | 0: T -1 , ˆ Σ 0: T -1 | 0: T -1 ) ) α 1: T = Dynamic Net ( x 0: T -1 | 0: T -1 ) Obtain ˆ F 1: T and ˆ H 1: T by (5) (ˆ µ 1: T | 0: T -1 , ˆ Σ 1: T | 0: T -1 ) = Prediction (ˆ µ 0: T -1 | 0: T -1 , ˆ Σ 0: T -1 | 0: T -1 , ˆ F 1: T ) ( w 1: T , r 1: T ) = encoder ( o 1: T ) ˆ K 1: T = ˆ Σ 1: T | 0: T -1 ˆ H 1: T M 1: T M T 1: T , M 1: T = GRU KG ( Conv ( ˆ Σ 1: T | 0: T -1 ) , r 1: T ) ˆ J 1: T = ˆ Σ 1: T | 1: T ˆ F T 1: T N 1: T N T 1: T , N 1: T = GRU SG ( Conv ( ˆ Σ 1: T | 0: T -1 ) ) (ˆ µ 1: T | 1: T , ˆ Σ 1: T | 1: T ) = Filtering (ˆ µ 1: T | 0: T -1 , ˆ Σ 1: T | 0: T -1 , ˆ K 1: T , w 1: T , ˆ H 1: T ) (ˆ µ 1: T | T , ˆ Σ 1: T | T ) = Smoothing (ˆ µ 1: T | 1: T , ˆ Σ 1: T | 1: T , ˆ J 1: T , ˆ F 1: T ) if (image imputation) then o 1: T = decoder (ˆ µ 1: T | T , ˆ Σ 1: T | T ) L ( o 1: T ) = Lower bound likelihood ( o 1: T ) by using (10) end if if (state estimation) then s 1: T = decoder (ˆ µ 1: T | T , ˆ Σ 1: T | T ) L ( s 1: T | o 1: T ) = Lower bound likelihood ( s 1: T ) by using (11) end if Backward Propagation by updating GRU cells with the steps in Theorem 3 or Theorem 4 .

To demonstrate the simplicity of our proposed GIN, we include intuitive inference code with Tensorflow library. The code runs with Python 3.6+. The entire code to reproduce the experiments are available in Github repository.

A.11.1 Python intuitive code.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims we have made in abstract and introduction are elaborated and achieved in the main paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations and assumptions necessary for Theorems 3 and 4. In the appendix, where we provide the proofs, we reiterate the requirements that the algorithm must meet.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The proofs of the mentioned theorems and mathematical claims are all provided in appendix section. In the main text, we refer the readers to the proofs after introducing each theorem-claim.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: All the required information for running the experiments are provided in the supplementary materials and anonymous github repository.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide instructions to synthesis the data and run the experiments. All details are provided in anonymous github repository.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide a detailed section in appendix to argue about network architecture, hyper parameter optimization, experiment setting, etc.. This addresses the question you asked.

Guidelines:

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We include confidence interval in our numerical experiments to show the reliability of the results. Initialization approaches are explained with details in the main script.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide detailed configuration of our machines for running the experiments.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We are following all the ethics provided by NeurIPS.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: We believe this work does not have negative impact on the social communities.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: The paper does not use existing assets.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification:

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a deep state-space model architecture with non-linear transitions and emissions. The model disentangles the latent representation for the dynamics and the one for the observed data at each time step - allowing therefore effective state estimation at future time steps and the ability to deal with missing data imputation.\nInference is performed with a deep Extended Kalman Filter, that relies on a RNN architecture to make a more efficient approximate computation of the Kalman Gain (KG) and smoothing gain (SG).\nThe method is tested on a number of simulated and realistic approaches, and it outperforms competing architectures.\n\nSTRENGTHS:\n1. Non-linear/deep state-space models are being used more and more in many applications. Parameter learning and state estimation is however challenging in this setting, and this paper provides an interesting method for this\n2. The method is more scalable than comptering KF-based methods thanks to the dynamics network approximation, but still effective despite the aproximation\n3. The method builds on some models in the literature, but provides some useful novel components\n4. The authors did extensive and well-though experiments/ablations, comparing with many SOTA models\n5. There is an extensive appendix covering many details that did not fit in the main text. I particularly appreciated \"A.11.1 Python intuitive code.\"\n\nWEAKNESSES:\nThe paper is not straightforward to read (had to read it carefully twice), mostly because of the way the required derivations are presented.\nThe notation used is somewhat not conventional within the ML-heavy NeurIPS community, and should be improved/clarified:\n1. In Section 4 you use o^+ notation which is not common in the ML community. Can you clarify what it means and why you need it? This explanation needs to be done in section 4, not referring to a different section.\n2. Similarly, what is s in line 219 and why do you need to introduce this notation? The sentence in line 219-221 is key but unclear\n3. Why do you need to define the \"gt\" in line 229?\n4. Not sure the SI perspective helps in the ML-heavy NeurIPS community, it brings confusion. Maybe can be added in the appendix?\n\nIn terms of novelty, the final model seems to me more similar to the Kalman VAE (KVAE) model than what the authors claim. Your model can almost be seen as a modification/extension to the KVAE in which you add the RNN approximation to avoid the O(n^3) complexity, model the transition noise covariance and use a slightly different parameterization for the dynamics network.\nCan you clarify the differences between your model and the KVAE? \nIn line 96 and Table 6 in the appendix, I don't think your KVAE description is correct: it has a setup very similar to your model which allows to estimate the state dynamics, and allows for direct optimization unlike what you claim.\nI am not as familiar with the other KF-based methods mentioned, but make sure your description is correct.\n\nMinor comments:\n1. Line 53: typo \"To to\"\n2. Line 71: typo \"We\" -> \"we\"\n3. You introduce gamma in line 144, but only say what it is in line 150, making the reader wonder if it was defined above and look for it\n4. Line 283: you say \"with n=3m\" without specifying what n and m are. Even if they are defined before, being a notation-heavy paper, better to remind the reader what n and m are.\n\nQUESTIONS:\nI am happy to increase the score as long as the comments/questions in the weaknesses section are clarified\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a very well theoretically motivated State-Space Model learning approach, which is implemented by a gated inference network. The network implements a Hammerstein-Wiener model within a modularized deep learning architecture. It uses GRU cells to mimic Kalman Filtering operations. Forward as well as forward-inverse processing routines optimize the hidden state space estimations. Several theoretical components add to the paper contribution. Evaluations show superior performance on several challenging toy problems with noisy data (pendulum, double pendulum, ball bouncing in irregular polygon, as well as odometry prediction from kiti data) evaluating state estimation and imputation tasks.\n\nSTRENGTHS:\nPaper is very well-structured. The work is also very well-motivated and well-embedded into the literature. \n\nThe theoretical motivation and system derivations are impressive and usefully embed the author’s GIN system into the Kalman Filtering background. Approximating everything in a variational inference manner via estimations of Gaussians and their Covariances is efficient.\n\nTheorems 3 and 4 offer a theoretical derivation for ensuring stability of the unfolding recurrence. \n\nThe evaluations contain sufficiently challenging problems. Performance is compared with many alternatives, showing superior performance nearly throughout. Only in Table 3 GIN was partially beaten by DeepVO.\n\nWEAKNESSES:\nThe theorems 3 and 4 are not really experimentally evaluated. Is instability observed when the recurrent matrix is not modified as proposed? The theorem’s proposal should be verified experimentally. \n\nEven more elaborate evaluations would of course be great. Seeing the great content and the importance of the theoretical derivations, though, I consider this a very minor point, which can be tackled in subsequent work.\n\nQUESTIONS:\nFigure 1 trajectories are not quite as smooth as I had expected. Any reason for this? The task, particularly between bounces, to predict a linear trajectory should be very simple. \n\nI do not really see how Theorems 1 and 2 are actual theorem. Don’t they just define the log likelihoods of the ground truth states / observations? \n\nTheorem 2 – is the second summand starting with factor $(1-o_t^{(k)})$ necessary? \n\nLine 260: should read \"in theorem 3\", no?\n\nLIMITATIONS:\nFurther evaluations and ablations to truly identify the core components that yield the great results of the GIN system.\n\nConfirm theory components experimentally.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-11"}

comments: SUMMARY:
This paper advances temporal reasoning in dynamic, high-dimensional, noisy environments by introducing a novel architecture for latent variable state space models. The architecture permits efficient Bayesian inference with nonlinear transitions and emissions. Experiments are performed on toy datasets and a simple real-world dataset for state estimation and missing data imputation, showing that it beats benchmarks relative to competing models like RNNs, autoregressive models, and latent variable approaches.

STRENGTHS:
Clear exposition of model architecture and inference algorithm. Theoretical analysis in Section 6.

WEAKNESSES:
I think one thing that could really strengthen this paper is showing an experiment on a more challenging data set / problem. The first two experiments are on toy problems.

I think another thing is to explain more clearly how this architecture is differentiated from others, i.e. the technical novelty. E.g. what is the relation of your model to other SSMs incorporating RNNs like the Variational RNN (which you benchmark against in the experiments), and what is it about that change that improves inference?

QUESTIONS:
Are there any experiments you could do that show a sequential modeling problem in which inference was previously intractable is now so?

LIMITATIONS:
Yes, limitations adequately addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-16