arXiv CSopen access

Stability Enhanced Gaussian Process Variational Autoencoders

Stability Enhanced Gaussian Process Variational Autoencoders

arXiv:2604.09331v1 [cs.LG] 10 Apr 2026

Carl R. Richardson∗

Jichen Zhang∗

Abstract— A novel stability-enhanced Gaussian process variational autoencoder (SEGP-VAE) is proposed for indirectly training a low-dimensional linear time invariant (LTI) system, using high-dimensional video data. The mean and covariance function of the novel SEGP prior are derived from the definition of an LTI system, enabling the SEGP to capture the indirectly observed latent process using a combined probabilistic and interpretable physical model. The search space of LTI parameters is restricted to the set of semi-contracting systems via a complete and unconstrained parametrisation. As a result, the SEGP-VAE can be trained using unconstrained optimisation algorithms. Furthermore, this parametrisation prevents numerical issues caused by the presence of a non-Hurwitz state matrix. A case study applies SEGP-VAE to a dataset containing videos of spiralling particles. This highlights the benefits of the approach and the application-specific design choices that enabled accurate latent state predictions.

I. I NTRODUCTION Variational autoencoders (VAEs) have proven to be very successful models for learning compact representations of high-dimensional data, in an unsupervised manner [1]. They have successfully been applied to static and time-varying tasks, including image classification [2], image segmentation [3], anomaly detection [4], long-horizon prediction [5], mobile robots [6], and data generation for autonomous driving [7]. In each of the time-varying applications, the latent process is subject to certain physical rules. Assuming knowledge of these rules is the foundation of model-driven approaches; however, these rules are often not completely known. Datadriven approaches, such as VAEs, typically assume no prior information but result in data and computationally expensive algorithms which lack interpretability. Koopman operator theory provides a mathematical framework for representing nonlinear dynamical systems as a linear, infinite-dimensional operator acting on observable functions of the system’s state [8]. This perspective enables the use of modern computational methods to approximate and analyse otherwise intractable systems [8], [9], [10]. Whilst the Koopman operator is theoretically required to be infinite-dimensional, for many practical applications, a finite-dimensional approximation is sufficient. This suggests CRR, JZ are with Department of Engineering Science, University of Oxford, Parks Road, Oxford, OX1 3PJ, UK {carl.richardson, jichen.zhang}@eng.ox.ac.uk EK is with Pacific Northwest National Laboratory, Richland, Washington, 99354, USA [email protected] JD is with Johns Hopkins University, Baltimore, Maryland, 21218, USA [email protected]. ∗ Equal contribution. This research was supported by the AT SCALE initiative via the Laboratory Directed Research and Development (LDRD) investments at Pacific Northwest National Laboratory (PNNL). PNNL is a national laboratory operated for the U.S. Department of Energy (DOE) by Battelle Memorial Institute under Contract No. DE-AC05-76RL0-1830.

Ethan King

Ján Drgoňa

VAEs coupled with linear latent dynamics (e.g., [11]) are also capable of approximating nonlinear latent dynamics. In this article, we assume the existence of a set of videos generated as a function of a low-dimensional latent process. The goal is to learn a model of the latent process from the video data. This setup is common in material science where, for example, a thin film deposition process cannot be directly measured, but video recordings of related diffraction patterns are available [12]. This problem has been studied through the lens of state space models [13], [14], [15], and Gaussian process (GP) based VAEs [11], [16], [17], [18]. Furthermore, to address the brittleness of many machine learning (ML) approaches for scientific problems, physicsinformed learning has emerged. This involves verifying properties such as stability after training [19], [20], [21], building prior knowledge into the model architecture [11], [22], [23], or encouraging desirable properties through the loss function design [24], [25]. Encoding prior knowledge into models reduces generality but improves performance; when the encoded structure is ubiquitous across applications, this trade-off preserves broad applicability while enabling the benefits of physics-informed learning [26]. One property common to many dynamical systems is stability, as highlighted by examples in physics [27], biology [28] and neural dynamics [29], [30]. Semicontraction is one definition of global stability [31] which implies the distance between two trajectories will never increase, regardless of their initial conditions. This is a robust form of stability [32], [33] and includes exponential stability and contraction as special cases. Furthermore, it can be applied to time-varying dynamical systems, such as those with external inputs. Contribution: We proposed a novel stability-enhanced Gaussian process (SEGP) where the mean and covariance functions are derived from the definition of a linear timeinvariant (LTI) system. We embed stability in the form of semi-contraction in the GP through an unconstrained and complete parametrisation (i.e., any LTI semi-contracting system can be constructed from the unconstrained variables). This parametrisation inherently prevents numerical issues which could arise during training due to a non-Hurwitz A matrix. The SEGP was incorporated within an existing GP-VAE framework to form the SEGP-VAE. This extends the physics-enhanced GP-VAE [11] by requiring less prior knowledge of the latent process and accounting for unknown initial conditions. The SEGP-VAE was trained on videos of a particle spiralling in a plane. After unsupervised training, the SEGP-VAE captures the indirectly observed underlying latent process using a combined probabilistic and interpretable physical model. When conditioned upon an observed video,

the SEGP posterior can accurately predict the underlying dynamics with low uncertainty. The LTI model can also be used for control design. Paper structure: The paper is organised such that the theoretical contribution is separated from the applicationspecific design choices. The background material and problem setup are presented in Section II and Section III. The SEGP is detailed in Section IV followed by the SEGP-VAE in Section V. Finally, a case study is presented in Section VI. This describes the data, application specific design choices such as encoder architecture, and presents empirical results. For other applications, these choices may differ. II. P RELIMINARIES A. Notation The set of real numbers greater than or equal to zero is denoted by R≥0 . The following sets of j × j matrices: symmetric (diagonal) positive definite, skew-symmetric and lower triangular with non-negative (positive) diagonal entries, are respectively denoted by Sj+ (Dj+ ), Skew(j) and Lj≥0 (Lj+ ). The matrix H with elements hjq is denoted by H = [hjq ]. A negative (semi-) definite matrix inequality is denoted by H ≺ 0 (H ⪯ 0). A vector function z(t) ∈ Rm at a set of discrete time points, T = {tj }N j=1 , is denoted by z(T ) = [z1 (t1 ) . . . z1 (tN ) . . . zm (t1 ) . . . zm (tN )] ∈ RmN . Similarly, a matrix function H(t, t′ ) = [hjq (t, t′ )] ∈ ∗ Rm×m over T and T∗ = {τq }N q=1 is denoted by

is given by      µy (T ) ŷ(T ) Ky (T, T ) + Σ ∼N , y(T∗ ) µy (T∗ ) Ky (T∗ , T )

! Ky (T, T∗ ) Ky (T∗ , T∗ ) (1) where Σ := Σ(T, T ) and Ky (T, T∗ ) = Ky (T∗ , T )⊤ . The choice of mean and covariance functions impose a modelling bias or an opportunity to encode prior knowledge into the model. For example, the popular squared exponential covariance function implies that the covariance between observations exponentially decays as a function of their separation in time. Refer to [35] for a comparison of different covariance functions. Assume the existence of a dataset Dŷ := (T, Ŷ), where Nŷ Ŷ := {ŷj (T )}j=1 is a batch of corrupted latent state trajectories. In this setting, y(t) ∼ GP(µy (t), Ky (t, t′ )) is referred to as the GP prior of the latent state. By Bayes’ Theorem, the predictive distribution of the latent state at the unobserved T∗ can be directly extracted from (1). The predictive distribution is given by (2)-(4) and provides a distribution conditioned on the dataset. The predictive distribution is referred to as the posterior when T∗ = T .
y(T∗ )|T∗ , Dŷ ∼ N µy|ŷ (T∗ ), Ky|ŷ (T∗ , T∗ ) (2) µy|ŷ (T∗ ) := µy (T∗ )+ Ky (T∗ , T ) Ky (T, T ) + Σ

−1

ŷ(T ) − µy (T )

(3)

−1

(4)

Ky|ŷ (T∗ , T∗ ) := Ky (T∗ , T∗ )− 

 H11 (T, T∗ ) . . . H1m (T, T∗ )   .. .. mN ×mN∗ .. H(T, T∗ ) = ∈ R . . . Hm1 (T, T∗ ) . . . Hmm (T, T∗ ) where Hil (T, T∗ ) = [hil (tj , τq )] ∈ RN ×N∗ . An m-dimensional vector-valued GP is denoted by z(t) ∼ GP(µz (t), Kz (t, t′ )) where µz (t) ∈ Rm is the vector-valued mean function and Kz (t, t′ ) ∈ Rm×m is the matrix-valued covariance function. The associated Gaussian distribution
over T is denoted by z(T ) ∼ N µz (T ), Kz (T, T ) where µz (T ) ∈ RmN and Kz (T, T ) ∈ RmN ×mN denote the mean vector and covariance matrix over T , respectively. B. Gaussian Processes This section provides a condensed background on GPs for vector-valued functions; refer to [34], [35] for a more comprehensive introduction. Two critical assumptions are made in the definition of a GP: (i) all variables can be represented as a sample from a multivariate Gaussian; (ii) any subset of variables can also be represented as a sample from a multivariate Gaussian. Suppose that i.i.d Gaussian noise is defined by η(t) ∼ GP(0, Σ(t, t′ )), with Σ(t, t′ ) = 2 diag(σ12 , . . . , σm ), and the latent state is governed by y(t) ∼ GP(µy (t), Ky (t, t′ )), then the corrupted latent state, ŷ(t) = y(t) + η(t), is ŷ(t) ∼ GP(µy (t), Ky (t, t′ ) + Σ(t, t′ )). By definition of a GP, the joint Gaussian of the corrupted latent state (observed at T ) and the latent state (at test points T∗ )

Ky (T∗ , T ) Ky (T, T ) + Σ

Ky (T, T∗ )

C. Latent Force Models Latent force models (LFM) were proposed in [36] for incorporating differential equations into GPs. The LFM framework can be applied to LTI systems (6) by noting the existence of an analytical solution of the form (5) where G : R+ × R+ → Rn×p is known as the Green’s function. Since a GP is closed under linear operators [34], placing a GP on u(t) ∼ GP(µu (t), Ku (t, t′ )) and assuming x(0) ∼ N (µx0 , Σx0 ) results in a GP x(t) ∼ GP(µx (t), Kx (t, t′ )) where µx (t) and Kx (t, t′ ) are determined by (5). Furthermore, since y is a linear combination of x and u, it is straightforward to determine y(t) ∼ GP(µy (t), Ky (t, t′ )). Z t A(t−t̃) (5) x(t) = eAt x(0) + |e {z B} u(t̃)dt̃ 0

:=G(t,t̃)

III. P ROBLEM S ETUP We consider the problem of modelling an unobserved latent process from a batch of related video recordings, N V V := {vj (T )}N j=1 , where T := {ti }i=1 . The pixel values of the square d-dimensional frames, over T , are denoted by vj (T ) = [v11 (t1 ) . . . v11 (tN ) . . . vdd (t1 ) . . . vdd (tN )]. As the latent states are not measured, this is an unsupervised learning problem with dataset D := (T, V). In this work, the latent process is assumed to be an LTI system defined by

(6) with parameters A ∈ Rn×n , B ∈ Rn×p , C ∈ Rm×n , D ∈ Rm×p , and initial condition x(0) ∼ N (µx0 , Σx0 ). ẋ(t) = A x(t) + B u(t) y(t) = C x(t) + D u(t)

(6a) (6b)

Several mild assumptions are made, these are: (i) u(t) ∼ GP(µu (t), Ku (t, t′ )) with known mean and covariance functions; (ii) the LTI parameters form a semi-contracting system (Theorem 2); (iii) the latent state is corrupted by i.i.d Gaussian noise; (iv) each frame of the video is only dependent upon the corrupted latent state, at the corresponding time.

Subbing in (8) and (9) leads to Ky (t, t′ ) = E[Y1 Y3⊤ ] + E[Y1 Y4⊤ ] + E[Y2 Y3⊤ ] + E[Y2 Y4⊤ ]
Y1 := C eAt x(0) − µx0 Z t

Y2 := GC (t, t̃) u(t̃) − µu (t̃) dt̃ + D u(t) − µu (t) 0
′ Y3 := C eAt x(0) − µx0 Z t′

Y4 := GC (t′ , t̂) u(t̂) − µu (t̂) dt̂ + D u(t′ ) − µu (t′ ) 0

We now focus on each term individually. ′

IV. S TABILITY E NHANCED G AUSSIAN P ROCESS

E[Y1 Y3⊤ ] = C eAt Σx0 eAt ⊤ C ⊤

This section derives the mean and covariance functions of the SEGP. This is later included within the SEGP-VAE, as in Fig. 1, to model the prior over the latent process. Theorem 1 derives the mean and covariance functions of a general LTI system (6). The matrices A, B, C, D are parameters of the GP which must be trained. Theorem 3 provides an unconstrained and complete parametrisation of the A matrix such that only semi-contracting LTI models can be learnt. Coupling Theorem 1 with the parametrisation in Theorem 3 results in the SEGP.
Theorem 1: Given u(t) ∼ GP µu (t), Ku (t, t′ ) and x(0) ∼ N (µx0 , Σx0 ), the latent state governed by (6) has the prior distribution y(t) ∼ GP(µy (t), Ky (t, t′ )) with the following covariance and mean functions ′

+

⊤

µy (t) = C e

x(0) − µx0

⊤  ⊤ u(t′ ) − µu (t′ ) D =0 {z }

covariance between x(0) and u(t′ )

as x(0) and u(·) are independent. By the same logic E[Y2 Y3⊤ ] = 0 Z t Z t′ Z 0t

GC (t, t̃)Ku (t̃, t̂)GC (t′ , t̂)⊤ dt̂dt̃

GC (t, t̃)Ku (t̃, t′ )D⊤ dt̃

Z t′ +

Ku (t, t̂)GC (t′ , t̂)⊤ dt̂

0

+ DKu (t, t′ )D⊤

DKu (t, t̂)GC (t , t̂) dt̂

(7)

GC (t, t̃) µu (t̃)dt̃ + D µu (t) 0 A(t̃−t̂)

(8)

where GC (t̃, t̂) := C e B is the Green’s function of (6) multiplied by C. Proof : The solution to (6a) is given by (5). Subbing this into (6b) gives Z t y(t) = C eAt x(0) + GC (t, t̃) u(t̃)dt̃ + D u(t) (9) 0

The LFM framework, summarised in Section II-C, highlights there must be a GP over y(t). The mean is derived as follows µy (t) = E[y(t)] = C eAt E[x(0)] +



0

Z t µx0 +

E |

+

0 At

+Ce

0

0 ′

covariance between x(0) and u(t̂)

At

E[Y2 Y4⊤ ] =

Ky (t, t′ ) = C eAt Σx0 eAt ⊤ C ⊤ + DKu (t, t′ )D⊤ Z t Z t′ + GC (t, t̃)Ku (t̃, t̂)GC (t′ , t̂)⊤ dt̂dt̃ 0 0 Z t + GC (t, t̃)Ku (t̃, t′ )D⊤ dt̃ Z t′

E[Y1 Y4⊤ ] = C eAt × Z t′ 

⊤  GC (t′ , t̂)⊤ dt̂ E x(0) − µx0 u(t̂) − µu (t̂) {z } 0 |

Z t GC (t, t̃) E[u(t̃)]dt̃ + D E[u(t)] 0

Subbing in the mean of x(0) and the mean function of u(·) results in (8). The covariance function between y(t) and y(t′ ) is defined by 

⊤  Ky (t, t′ ) = E y(t) − µy (t) y(t′ ) − µy (t′ )

Subbing each term into Ky (t, t′ ) results in (7). □ In [11] the authors presented a covariance function for (6)
when D = 0, x(0) = 0 and u(t) ∼ GP 0, Ku (t, t′ ) . We extend this to widen the scope of applicability to account for the cases when D ̸= 0 (e.g., analogue circuits), x(0) ∼ N (µx0 , Σx0 ), such as aircraft tracking, and an input signal with non-zero mean function (e.g., stochastic interest rates). Directly applying Theorem 1 is problematic because nonHurwitz A matrices are permitted, which destabilise training due to unbounded mean and covariance functions. Assumption (ii) in Section III address this issue; however, assuming the LTI parameters form a semi-contracting system requires the following well established result to hold. Theorem 2 (Semi-contracting LTI system): The LTI system (6) is semi-contracting if and only if there exists P ∈ Sn+ which satisfies the following linear matrix inequality (LMI) P A + A⊤ P ⪯ 0 (10) Proof : Refer to the appendix. □ The parameters of (6) that satisfy Theorem 2 are defined by the set Ω := {(A, B, C, D) | ∃P ∈ Sn+ : P A + A⊤ P ⪯ 0}

(11)

where B, C, D are unconstrained. The following result provides a complete parametrisation of the set Ω in terms of unconstrained variables, allowing the SEGP to be implemented using parametrised layers and trained using unconstrained optimisers, like gradient descent. Theorem 3: (A, B, C, D) ∈ Ω if and only if there exists V1 ∈ Ln+ , V2 ∈ Ln≥0 and V3 ∈ Skew(n) related to A by 1 −1 A = − P −1 V2 V⊤ V3 2 +P 2

P = V1 V⊤ 1

P (v, y) = P (v | y)Pψ (y) =

= (12)

with B ∈ Rn×p , C ∈ Rm×n and D ∈ Rm×p . Proof (Necessity): If (A, B, C, D) ∈ Ω, there exists P ∈ Sn+ such that P A + A⊤ P ⪯ 0. As the left hand side is symmetric and negative semi-definite, a Cholesky decomposition must exist. That is P A + A⊤ P = − V 2 V ⊤ 2

distribution can be expressed as the product in (14b), where Pψ (y) denotes the SEGP prior with trainable parameters ψ.

(13)

This equality is satisfied if and only if 1 P A = − V2 V⊤ 2 + V3 2 where V3 ∈ Skew(n). Reconstructing A results in (12). (Sufficiency): If there exists V1 as stated in Theorem 3, it immediately follows that P is positive definite, symmetric and full rank. Next, constructing A as in (12) gives 1 1 ⊤ ⊤ P A + A⊤ P = − V 2 V ⊤ 2 + V3 − V2 V2 + V3 2 2 = − V2 V⊤ 2 which is negative semi-definite, symmetric and equal to the decomposition in (13). □ As specified in Theorem 3, the SEGP has trainable parameters ψ := {V1 , V2 , V3 , B, C, D}. The resulting space of A matrices unconditionally satisfies Theorem 2, ensuring (A, B, C, D) ∈ Ω throughout training. This ensures trajectories generated by the SEGP will be bounded, since any A matrix satisfying (10) has eigenvalues with non-positive real components [37]. Thus numerical issues due to the term eAt are avoided when computing the mean and covariance. Furthermore, the sets Ln≥0 , Ln+ can be considered as unconstrained since any positive element can be expressed as the the magnitude of a scalar plus a small positive value. An analogous case holds for skew-symmetric matrices. Finally, as diagonal elements of triangular matrices correspond to their eigenvalues [38], P will always be full rank, ensuring the existence of it’s inverse. V. S TABILITY E NHANCED G AUSSIAN P ROCESS VARIATIONAL AUTOENCODER This section describes the general form of the SEGP-VAE architecture and the unsupervised training objective. Consider the joint distribution between the video and the latent state, where v := v(T ), y := y(T ), vi := v(ti ), and yi := y(ti ). The joint distribution can be expressed as (14a) since the video is conditionally dependent on the latent state. By assumption (iv) in Section III, the conditional

N Y i=1 N Y

(14a)

P (vi | yi )Pψ (y)

(14b)

B vi | pθ (yi ) Pψ (y)

(14c)

i=1

Finally, in this work we make the simplifying assumption that 2 the pixel values are black and white (i.e., vi ∈ {0, 1}d ). This allows the Bernoulli to be used to model the likelihood of the video, as in (14c). The Bernoulli is parametrised by 2 pθ : Rm → [0, 1]d which represents the probability of each pixel being white. This function is implemented by the decoder of the VAE, with trainable parameters θ. Any universal function approximator may be used for the decoder providing each output is restricted to [0, 1]. For example, an MLP with sigmoid output activations. Note that the Bernoulli could easily be replaced by the Beta distribution when modelling video with normalised RGB pixel values. In this case, the decoder would map the latent state to the shape parameters of the Beta distribution. Due to the Bernoulli distribution (or Beta), there exists no analytical solution for the posterior, P (y | v); hence, the variational approximation of the posterior proposed in [18] is reused N Y

1 N yi | µϕ (vi ), Σϕ (vi ) Pψ (y) Qϕ,ψ y | v = Lϕ,ψ (v) i=1 | {z } :=qϕ (yi | vi )

(15) where Lϕ,ψ (v) denotes the marginal likelihood. This approximation is simply the true posterior with only the troublesome likelihood replaced by a Gaussian. The mean and 2 covariance functions of the likelihood, µϕ : {0, 1}d → Rm , 2 Σϕ : {0, 1}d → Dm + respectively, are implemented by the VAE encoder, with trainable parameters ϕ. Any universal function approximator may be used for the encoder providing the outputs, corresponding to the diagonal elements of Σϕ , are guaranteed to be positive. For example, mapping these outputs through an exponential function guarantees this. The SEGP-VAE architecture is depicted in Fig. 1. We consider how to learn an accurate approximation of the posterior in Section VI-C; for now, we just highlight how Qϕ,ψ (y | v) is computed. Consider a single video, v(T ), from the dataset. Mapping it through the encoder results in an associated set of latent “observations”, L := {µϕ (vi )}N i=1 , each with noise Σϕ (vi ). Conditioning the GP prior on these observations yields the analytically tractable posterior (15) which could be extracted from a joint Gaussian of the form (1) where y(T ) and Σ are replaced by µϕ (v) ∈ RmN and Σϕ (v) ∈ DmN + , as defined below   µϕ (v) = µϕ,1 (v1 ). . . µϕ,1 (vN ). . . µϕ,m (v1 ). . . µϕ,m (vN )
2 2 2 2 (v1 ). . . σϕ,1 (vN ). . . σϕ,m (v1 ). . . σϕ,m (vN ) Σϕ (v) = diag σϕ,1 As a result, the mean and covariance of the posterior can be directly computed from the standard GP equations (2)-(4).

SEGP-VAE Latent state trajectories Unknown excitation

Encoder

Decoder Time SEGP

Input Fig. 1.

A. Training As this is an unsupervised learning task, the goal was to obtain an accurate reconstruction of the input video, whilst indirectly training the SEGP to accurately model the latent process. To this end, an augmented version of the evidence lower bound (ELBO) objective was employed. The ELBO objective is defined as LELBO (·) = EQϕ,ψ

N hX

i log B vi | pθ (yi )

i=1

+ β · DKL Pψ (y)||Qϕ,ψ (y | v)

(16)

where the expectations are with respect to the variational posterior (15). The first term is the reconstruction error which indirectly depends on the encoder, ϕ, and SEGP, ψ, as yi ∼ Qϕ,ψ (yi | vi ) is sampled using the reparameterisation trick [39]. The second term is the analytically tractable Kullback–Leibler (KL) divergence, which regulates how far the variational posterior deviates from the SEGP prior. Typically, the weighting parameter β is set to one; however, we treat this as a hyperparameter to be tuned. We augment the ELBO loss by introducing an additional regularisation term on the SEGP parameters, r(ψ), as shown below. This term is also weighted by a tunable hyperparameter, λ. The choice of regularization term is task dependent and enables training of the SEGP to be biased. L(ϕ, ψ, θ; v) = LELBO (ϕ, ψ, θ; v) + λr(ψ)

(17)

VI. C ASE S TUDY: S PIRALLING PARTICLE Prior to this section, the general SEGP-VAE model was presented; however, as with all ML approaches, some problem specific choices must be made to achieve strong performance. This section begins by detailing the data which is subsequently used to justify certain choices in the model and training setup. Finally, the empirical results are presented. All code and weights are freely available1 . A. Data A dynamical system of the form (6), where x = [r θ]′ , represents the position of a particle in a plane in terms of its polar coordinates (radius, r, and angle from the positive x-axis, θ). The following matrices were used     −0.6 0 0 Â = B̂ = Ĉ = I D̂ = 0 0 0 1 1 https://github.com/pnnl/COPIP

Output

SEGP-VAE architecture with semi-contracting kernel.

with the external input u(t) ∼ GP µu (t), ku (t, t′ ) deter
mined by µu (t) = 0.4πt and ku (t, t′ ) = exp −0.5(t−t′ )2 . 2 Initial conditions were sampled from x(0) ∼ N (µx0 , σx0 I) ′ where µx0 = [1.5 0] and σx0 = 0.2. Euler integration was used to simulate the trajectories over a 3s period with step size 1×10−2 . The trajectories were discretised with a sample period of Ts = 12×10−2 (N = 25) and additive measurement noise was applied to the discretised trajectories with Σ = σ 2 I and σ 2 = 1 × 10−3 . The noisy discretised trajectories were transformed to Cartesian coordinates, rescaled to pixel indices and rendered as a ball, with a radius of 2 pixels, onto a H × W binary pixel canvas, where H = W = 40. In total, Nv = 40, 000 videos were generated by sampling Nv external inputs and 1 initial condition for each. Example videos are shown in Fig. 5, where the images at each time step have been overlaid to show the full trajectory. B. SEGP-VAE 1) Encoder: Since we are working with images, the encoder of the SEGP-VAE was implemented by a CNNbased architecture. This consisted of three convolutional layers, with the first two having stride 2, each followed by Group Normalization and ReLU activations, and the last one having stride 1. This produced a set of latent feature maps F ∈ RK×H×W , where K = 1 was the chosen number of key points. Due to the rotational dynamics in the data, a two-step spatial-to-polar mapping was implemented: (i) For each channel, k, a Spatial Softmax layer [40] mapped the latent features, w, to specific image coordinates (xk , yk ). The expected coordinate in the normalized image space was computed as: E[vk ] =

H X W X h=1

exp(Fk,h,w /τ ) P · posv (w) i,j exp(Fk,i,j /τ ) w=1

where vk ∈ {xk , yk }, τ = 1.0 was the chosen value of the temperature parameter controlling the sharpness of the attention, and posv (w) denotes the normalized horizontal or vertical position. This operation is fully differentiable. (ii) The Cartesian coordinates were mapped to polar coordinates (r, θ) via the atan2 function. However, the inherent 2π-discontinuity of the phase component posed a significant challenge for gradient-based optimisation in temporal sequences. To resolve this, we implemented a temporal unwrapping trick. Given a sequence of raw angles {θt }N t=1 ,

By eliminating artificial jumps at ±π, the latent state zt = [rt , θt′ ]⊤ evolves in a continuous space. Finally, the encoder mean was obtained by applying an affine transformation to the polar coordinates with learnable scale and bias parameters. The encoder variance was derived from learnable logvariance parameters clamped for numerical stability. 2) Decoder: In contrast, the decoder was implemented by a simple 2-layer MLP. Batch norm was applied to the 500 hidden neurons, followed by ReLU activations. This was followed by a linear output layer, parametrising the Bernoulli likelihood. 3) SEGP: The SEGP was implemented according to Theorem 1, with the A matrix parametrised according to Theorem 3. To make the problem tractable, the matrices B, C, D were assumed to be known; hence ψ = {V1 , V2 , V3 }. This corresponds to having information of the input-output structure of the LTI system, but no knowledge of the internal dynamics. Whilst this isn’t a general requirement, we found structural identifiability to be an issue for this particular task. That is, coupled dynamics could be learnt without substantially affecting the reconstruction objective. C. Training Setup 1) Objective Function: To ensure that the learned latent dynamical system remains physically interpretable and avoids over-parameterization, we chose the the L1 regularization, on the state matrix A. Specifically, this led to r(ψ) = P ∥A∥1 in (17) where ∥A∥1 = i,j |Ai,j | denotes the entrywise L1 norm. The L1 regularization promotes sparsity in the state matrix and suppresses excessive cross-dimensional coupling. This acts as a complexity control mechanism, encouraging simpler and more interpretable dynamics. We also employed a linear scheduler for the weight of the L1 penalty term, λ. Over the course of training, this increased the weighting of this term from 0.025 to 0.3. Furthermore, we empirically found that setting β = 2.5 prevented the reconstruction error from excessively dominating the objective. 2) Feature Scaling: As shown in Fig. 6, the angular position can be much larger than the radius. Without intervention, this feature would naturally dominate the KL divergence. To ensure both features were treated equally during training, the KL divergence was computed in the normalised space. For the prior and posterior, this involved normalising each element of the mean function and covariance matrix by the corresponding largest radius or angular position from the training dataset. 3) Optimizer and Parameter Initialisation: All model parameters including encoder, decoder and SEGP were trained jointly using the AdamW optimizer with a learning rate of 5 × 10−3 and weight decay 10−5 . Kaiming initialisation was used for the encoder and decoder, whilst the A matrix (12) was initialised with V1 = I, V3 = 0, and V2 sampled from a zero-mean Gaussian with standard deviation 10−3 .

Per-Pixel Reconstruction Loss

we computed the unwrapped sequence {θt′ } such that:   t X θi−1 − θi θt′ = θt + 2π · round 2π i=2

Train Test

0.05 0.04 0.03 0.02 0.01 0

Fig. 2.

100

200

300 Epoch

400

500

600

Evolution of per-pixel reconstruction error during training. KL Divergence

L1 Norm

3 2 KL Train KL Test

1 0

200

400 Epoch

Fig. 3.

600

0

200

400

600

Epoch

KL divergence (unweighted) and L1 norm during training.

D. Empirical Results 1) Training: We report the results over 600 training epochs, with both training and testing curves monitored to evaluate effectiveness. Fig. 2 shows the ELBO and reconstruction curves decrease rapidly during early epochs and then smoothly converge to a stable equilibrium. Training and testing curves indicates good generalization and no overfitting. The KL divergence in Fig. 3 initially decreases slightly and then increases, eventually stabilizing around a steady non-zero value. This shows that posterior collapse is avoided and highlights that the posterior actively contributes to the accurate video reconstruction. The L1 norm in Fig. 3 increases modestly during early training before steadily decreasing. Whilst some terms in the objective function converge quickly, this regularisation term consistently varies over the entire training period. As this corresponds to consistent reduction in the loss function, it suggests this term plays an important role during training. 2) SEGP versus Squared Exponential GP: To see the benefit of the SEGP over a typical GP, this section compares the covariance matrix of the learnt SEGP prior with that of a standard multi-output GP, where each state was modelled by a zero-mean function and distinct, trainable squared exponential (SE) kernel. The standard GP was trained directly on the trajectory data, until convergence. From the parameters defined in Section VI-A, the uncertainty of the radius only depends on the initial condition. Furthermore, the radius dynamics are contracting which implies the uncertainty should exponentially disappear. This is reflected in the SEGP covariance of the radius (Fig. 4); however, it can be more clearly seen in Fig. 6 since the uncertainty of the radius is negligible compared to the angle.

10

True Prediction

5

0 0.5

1.0

1.5

0.5

1.0

1.5

0.5

1.0

1.5

0.5

1.0

1.5

0.5

1.0

1.5

Fig. 4. Learnt covariance matrices of SEGP prior (left) & SE kernel (right). Top left block is the covariance of the radius and bottom right block is the covariance of the angular position. Top right and bottom left is the covariance across dimensions.

On the other hand, the uncertainty of the angular position is independent of the initial conditions, but integrates the uncertainty of the input over time. Fig. 4 shows that the learnt SEGP covariance of the angle represents this whilst the covariance of the standard GP is restricted to being an exponentially decaying function of |t − t′ |. Finally, both covariance matrices show the radius and angular position are uncorrelated; however, this was enforced by the standard GP and learnt by the more flexible SEGP. 3) Posterior versus Prior: Fig. 7 compares the average absolute error between the posterior (and prior) mean prediction and the latent trajectories from the test set. Furthermore it compares the average variance of the posterior and prior across the test set. Only the angular position component is considered since there is such little variance in the radius component of the trajectory data. Since the average absolute error and variance of the posterior (conditioned on the videos) is significantly reduced compared to the prior, this highlights that the SEGP-VAE enocoder has learnt a representation which translates the observed evolution of the particle in the video into an appropriate Bayesian update. The low error and variance of the posterior are also highlighted in Fig. 6 where the posterior mean aligns tightly with the ground truth and the uncertainty bands are narrow. 4) Prediction: In Fig. 5 the reconstructed frames accurately reproduce the rotational motion present in the ground truth sequences. More importantly, in Fig. 5 and Fig. 6 the posterior mean closely matches the true latent trajectories with narrow uncertainty bands. This is more rigorously supported by Fig. 7 showing the average absolute error between the SEGP posterior mean and the latent trajectories in the test set. Furthermore, the error in the spectral norm between the learnt and ground truth A matrices was only ∥Â−A∥2 = 0.016; hence the worst case directional distortion of the learnt latent process was minimal. VII. C ONCLUSION A novel Stability Enhanced Gaussian Process (SEGP) was proposed with mean and covariance functions derived from the definition of a semi-contracting LTI system. This provided a probabilistic and interpretable physical modelling tool whilst avoiding numerical issues caused by the presence of a non-Hurwitz state matrix. The SEGP was used within

Fig. 5. Randomly sampled videos from test set (top); corresponding latent trajectories (radius along x-axis, angular position along y-axis) and the posterior mean prediction (middle); reconstructed video (bottom)).

a VAE framework for training the LTI system to model the latent process which generated the video data. Whilst the latent process in the case study was a linear system, we expect our approach to also be suitable for nonlinear systems based on Koopman theory. The case study highlighted: (i) the superior flexibility of the SEGP over a standard GP kernel; (ii) that strong performance required significant application specific design of the encoder and augmentation of the standard ELBO loss. Whilst not a theoretical requirement of the approach, prior knowledge of the underlying processes input–output structure was needed in practice. Overcoming this limitation will be the subject of future work. A PPENDIX The fundamental result of contraction analysis is presented in [31, Section 3]. For (6) to be semi-contracting, the distance between any two trajectories, in a Riemann space, must never increase. This distance is denoted by V (δx), where δx denotes the virtual displacement between two trajectories n ˙ . The virtual displacement is governed by δx. and P ∈ S+ V (δx) = δx⊤ P δx

˙ = Aδx δx

(18)

Proof of Theorem 2: Consider V̇ (δx) ˙ V̇ (δx) = 2δx⊤ P δx = 2δx⊤ P Aδx = δx⊤ (P A + A⊤ P )δx n If there exists a matrix P ∈ S+ satisfying (10), then V̇ (δx) ≤ 0 ∀ δx. Hence, the system is semi-contracting. □

R EFERENCES [1] C. Doersch, “Tutorial on variational autoencoders,” arXiv preprint arXiv:1606.05908, 2016. [2] K. Gregor, I. Danihelka, A. Graves, D. Rezende, and D. Wierstra, “Draw: A recurrent neural network for image generation,” in International conference on machine learning. PMLR, 2015, pp. 1462–1471. [3] 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. [4] F. Waseem, R. Martinez, and C. Wu, “Visual anomaly detection in video by variational autoencoder,” arXiv preprint arXiv:2203.03872, 2022.

1.5

10.0

True Prior 95% CI

1.0

True Prior 95% CI

7.5 5.0 2.5

0.5

0.0

1.5

True Prediction 95% CI

1.0

True Prediction 95% CI

6 4 2

0.5

0 0

1

2

3

0

1

Time (t)

2

3

Time (t)

5

Prior Posterior

1.5

Variance (Angle)

Absolute Error (Angle)

Fig. 6. Randomly sampled latent trajectory from test set, plotted against the learnt SEGP prior (top) and learnt SEGP posterior (bottom). Radius plotted on left and angular position on right.

1.0 0.5

Prior Posterior

4 3 2 1 0

0.0 0

1

Time (t)

2

3

0

1

2

3

Time (t)

Fig. 7. Absolute error between the learnt SEGP prior (posterior) mean and the angular position trajectory data, averaged over the test set (left). The variance of the learnt SEGP prior (posterior) for the angular position, averaged over the test set (right).

[5] V. Saxena, J. Ba, and D. Hafner, “Clockwork variational autoencoders,” Advances in Neural Information Processing Systems, vol. 34, pp. 29 246–29 257, 2021. [6] M. Nagano, T. Nakamura, T. Nagai, D. Mochihashi, and I. Kobayashi, “Spatio-temporal categorization for first-person-view videos using a convolutional variational autoencoder and gaussian processes,” Frontiers in Robotics and AI, vol. 9, p. 903450, 2022. [7] A. Amini, W. Schwarting, G. Rosman, B. Araki, S. Karaman, and D. Rus, “Variational autoencoder for end-to-end control of autonomous driving with novelty detection and training de-biasing,” in 2018 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). IEEE, 2018, pp. 568–575. [8] M. Budišić, R. Mohr, and I. Mezić, “Applied Koopmanism,” Chaos: An Interdisciplinary Journal of Nonlinear Science, vol. 22, no. 4, 2012. [9] Y. Lian and C. N. Jones, “On Gaussian process based Koopman operators,” IFAC-PapersOnLine, vol. 53, no. 2, pp. 449–455, 2020. [10] K. Miao, H. Wang, X. Ding, K. Gatsis, A. Krause, and A. Papachristodoulou, “Learning koopman representations with controllability guarantees,” in The Fourteenth International Conference on Learning Representations, 2026. [11] T. Beckers, Q. Wu, and G. J. Pappas, “Physics-enhanced Gaussian process variational autoencoder,” in Learning for Dynamics and Control Conference. PMLR, 2023, pp. 521–533. [12] T. C. Kaspar, S. Akers, H. W. Sprueill, A. H. Ter-Petrosyan, J. A. Bilbrey, D. Hopkins, A. Harilal, J. Christudasjustus, P. Gemperline, and R. B. Comes, “Machine-learning-enabled on-the-fly analysis of RHEED patterns during thin film deposition by molecular beam epitaxy,” Journal of Vacuum Science & Technology A, vol. 43, no. 3, 2025. [13] M. Fraccaro, S. Kamronn, U. Paquet, and O. Winther, “A disentangled recognition and nonlinear dynamics model for unsupervised learning,” Advances in neural information processing systems, vol. 30, 2017. [14] W. Lin, N. Hubacher, and M. Khan, “Variational message passing with structured inference networks,” arXiv preprint arXiv:1803.05589, 2018.

[15] M. Pearce, S. Chiappa, and U. Paquet, “Comparing interpretable inference models for videos of physical motion,” in 1st symposium on advances in approximate bayesian inference, 2018. [16] F. P. Casale, A. Dalca, L. Saglietti, J. Listgarten, and N. Fusi, “Gaussian process prior variational autoencoders,” Advances in neural information processing systems, vol. 31, 2018. [17] A. Campbell and P. Liò, “tvGP-VAE: Tensor-variate Gaussian process prior variational autoencoder,” arXiv preprint arXiv:2006.04788, 2020. [18] M. Pearce, “The Gaussian process prior VAE for interpretable latent dynamics from pixels,” in Symposium on advances in approximate bayesian inference. PMLR, 2020, pp. 1–12. [19] C. R. Richardson, M. C. Turner, and S. R. Gunn, “Strengthened Circle and Popov Criteria for the stability analysis of feedback systems with ReLU neural networks,” IEEE Control Systems Letters, 2023. [20] C. R. Richardson, M. C. Turner, S. R. Gunn, and R. Drummond, “Strengthened stability analysis of discrete-time Lurie systems involving ReLU neural networks,” in Learning for Decision and Control (L4DC). L4DC, 2024. [21] C. R. Richardson, M. C. Turner, and S. R. Gunn, “Analysis of lurie systems with magnitude nonlinearities and connections to neural network stability analysis,” IEEE Transactions on Automatic Control, 2026. [22] M. M. Bronstein, J. Bruna, T. Cohen, and P. Veličković, “Geometric deep learning: grids, groups, graphs, geodesics, and gauges,” arXiv preprint arXiv:2104.13478, 2021. [23] C. R. Richardson, M. C. Turner, and S. R. Gunn, “Lurie networks with robust convergent dynamics,” Transactions on Machine Learning Research, 2025. [24] J. Degrave, F. Felici, J. Buchli, M. Neunert, B. Tracey, F. Carpanese, T. Ewalds, R. Hafner, A. Abdolmaleki, D. de Las Casas et al., “Magnetic control of tokamak plasmas through deep reinforcement learning,” Nature, vol. 602, no. 7897, pp. 414–419, 2022. [25] E. King, Y. Li, S. Hu, and E. Machorro, “Physics-informed machinelearning model of temperature evolution under solid phase processes,” Computational Mechanics, vol. 72, no. 1, pp. 125–136, 2023. [26] J. Drgoňa, T. X. Nghiem, T. Beckers, M. Fazlyab, E. Mallada, C. Jones, D. Vrabie, S. L. Brunton, and R. Findeisen, “Safe physicsinformed machine learning for dynamics and control,” in 2025 American Control Conference (ACC), 2025, pp. 591–606. [27] H. K. Khalil, “Nonlinear systems,” Patience Hall, vol. 115, 2002. [28] L. Kozachkov, M. Lundqvist, J.-J. Slotine, and E. K. Miller, “Achieving stable dynamics in neural circuits,” PLoS computational biology, vol. 16, no. 8, p. e1007659, 2020. [29] H. B. Mohammadi, S. Hauberg, G. Arvanitidis, N. Figueroa, G. Neumann, and L. Rozo, “Neural contractive dynamical systems,” in The Twelfth International Conference on Learning Representations, 2024. [30] J. Drgoňa, A. Tuor, S. Vasisht, and D. Vrabie, “Dissipative deep neural dynamical systems,” IEEE Open Journal of Control Systems, vol. 1, pp. 100–112, 2022. [31] W. Lohmiller and J.-J. E. Slotine, “On contraction analysis for nonlinear systems,” Automatica, vol. 34, no. 6, pp. 683–696, 1998. [32] A. Davydov and F. Bullo, “Perspectives on contractivity in control, optimization, and learning,” arXiv preprint arXiv:2404.11707, 2024. [33] S. Jaffe, A. Davydov, D. Lapsekili, A. Singh, and F. Bullo, “Learning neural contracting dynamics: Extended linearization and global guarantees,” arXiv preprint arXiv:2402.08090, 2024. [34] C. K. Williams and C. E. Rasmussen, Gaussian processes for machine learning. MIT press Cambridge, MA, 2006, vol. 2, no. 3. [35] M. A. Alvarez, L. Rosasco, N. D. Lawrence et al., “Kernels for vectorvalued functions: A review,” Foundations and Trends® in Machine Learning, vol. 4, no. 3, pp. 195–266, 2012. [36] M. A. Alvarez, D. Luengo, and N. D. Lawrence, “Linear latent force models using Gaussian processes,” IEEE transactions on pattern analysis and machine intelligence, vol. 35, no. 11, pp. 2693–2705, 2013. [37] J. Hespanha, Linear systems theory. Princeton university press, 2018. [38] K. B. Petersen, M. S. Pedersen et al., “The matrix cookbook,” Technical University of Denmark, vol. 7, no. 15, p. 510, 2008. [39] D. P. Kingma and M. Welling, “Auto-encoding variational bayes,” arXiv preprint arXiv:1312.6114, 2013. [40] S. Levine, C. Finn, T. Darrell, and P. Abbeel, “End-to-end training of deep visuomotor policies,” Journal of Machine Learning Research, vol. 17, no. 39, pp. 1–40, 2016.