Discrete-Time Conditional Gaussian Koopman Network

• Discrete-Time CGKN is a learned surrogate that maps unobserved states to a conditional Gaussian latent space, enabling analytic data assimilation and efficient forecasting for nonlinear systems.
• It leverages a Koopman-inspired latent embedding to yield conditional linear dynamics, bypassing the need for expensive ensemble-based filtering.
• The model directly learns the coarse time-Δt solution operator, achieving competitive performance on PDE benchmarks while significantly reducing computational cost.

Discrete-Time Conditional Gaussian Koopman Network (CGKN) denotes a class of learned surrogate models for partially observed nonlinear dynamical systems in which the unobserved state is mapped to a latent Koopman-style representation whose evolution is linear conditional on the observed variables. The resulting joint model is a conditional Gaussian nonlinear system, so the latent posterior remains Gaussian and can be updated analytically for data assimilation (DA). In the literature, CGKN first appears as a continuous-time framework based on neural differential equations and is then specialized to a directly discrete-time formulation that learns the coarse time-$\Delta t$ solution map itself rather than relying on internal time integration (Chen et al., 2024, Chen et al., 11 Jul 2025).

1. Emergence and scope

The defining motivation of CGKN is the joint treatment of forecasting and DA within one trainable model. Conventional deep surrogates can approximate nonlinear dynamics well, but their black-box structure makes DA expensive and typically forces the use of ensemble-based methods with sampling error and empirical tuning. CGKN addresses this by retaining enough nonlinear structure to model complex dynamics while enforcing a conditional Gaussian latent architecture that yields analytic filtering formulae (Chen et al., 2024).

The transition from the 2024 formulation to the discrete-time version is not merely a change of notation. The continuous-time model is written as a conditional Gaussian neural SDE/ODE, whereas the discrete-time model learns the time-$\Delta t$ map directly at the temporal resolution of the data. This shift is especially important for PDE surrogates trained on coarse output intervals, because it removes the need for expensive numerical integration inside the learned model (Chen et al., 11 Jul 2025).

Formulation	Core latent structure	Main emphasis
Continuous-time CGKN	$$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$	Forecasting and analytic DA in neural differential equations
Discrete-time CGKN	$$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$	Direct learning of the coarse solution map
Lagrangian CGKN	Same conditional Gaussian latent form, with tracer-specific observation structure	DA and prediction from Lagrangian observations

A central conceptual point is that CGKN is Koopman-inspired rather than fully Koopman-linear. Standard Koopman learning typically seeks a lifted state with globally linear dynamics. CGKN instead learns a latent embedding of the hidden state and only requires linearity in that latent state conditional on the observed variables. This substantially enlarges the admissible model class while preserving tractable posterior inference (Chen et al., 2024, Chen et al., 11 Jul 2025).

2. Discrete-time state-space formulation

The discrete-time setting begins with a partially observed nonlinear dynamical system

$$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$

where $\mathbf{u}_1^n$ denotes observed variables and $\mathbf{u}_2^n$ denotes unobserved variables. In general,

$$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$

CGKN replaces this with a surrogate in which the hidden component is encoded into a latent variable $\mathbf{v}^n$ through an autoencoder: $$\mathbf{v}=\boldsymbol{\varphi}(\mathbf{u}_2), \qquad \mathbf{u}_2=\boldsymbol{\psi}(\mathbf{v}).$$ The learned surrogate is

$\Delta t$0

with independent standard Gaussian noises $\Delta t$1 (Chen et al., 11 Jul 2025).

This model is nonlinear in the full state because $\Delta t$2 depend nonlinearly on $\Delta t$3. Yet for fixed $\Delta t$4, the update is affine in $\Delta t$5. That is the precise sense in which the model is conditional Gaussian. It is also the sense in which the latent representation is “Koopman”: the latent dynamics are simple—linear in $\Delta t$6—once conditioned on the observed coordinates (Chen et al., 11 Jul 2025).

The continuous-time predecessor uses the analogous structure

$\Delta t$7

and the discrete-time model can be viewed as learning the corresponding time-discretized solution operator directly (Chen et al., 2024).

3. Conditional Gaussian inference and analytic data assimilation

Because the latent dynamics are linear-Gaussian once the observed trajectory is fixed, the posterior

$\Delta t$8

is Gaussian. Writing

$\Delta t$9

the discrete-time conditional Gaussian filter is

$$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$0

with

$$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$1

These recursions are Kalman-like but with coefficients depending on $$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$2 through the learned networks (Chen et al., 11 Jul 2025).

The physical hidden state is recovered through the decoder: $$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$3 This approximation is exact only if the decoder is linear; with a nonlinear decoder, $$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$4 is used as a practical posterior mean surrogate and is explicitly optimized during training through the DA loss (Chen et al., 11 Jul 2025).

A major computational consequence is that DA no longer requires ensembles. In the discrete-time benchmarks, the latent filter is analytic and significantly cheaper than EnKF applied to the full PDE model. For viscous Burgers’, the reported DA cost per trajectory is $$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$5 s for CGKN versus $$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$6 s for EnKF; for Kuramoto–Sivashinsky it is about $$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$7 s versus $$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$8 s; for $$\frac{d\mathbf{v}}{dt}=\mathbf{f}_2(\mathbf{u}_1)+\mathbf{g}_2(\mathbf{u}_1)\mathbf{v}+\boldsymbol{\sigma}_2\dot{\mathbf{W}}_2$$9-D Navier–Stokes it is about $$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$0 s versus about $$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$1 s on the same hardware (Chen et al., 11 Jul 2025).

4. Learning objectives, noise estimation, and implementation

The discrete-time formulation trains four coupled components: an encoder $$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$2, a decoder $$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$3, and a dynamics network producing $$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$4 as functions of $$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$5. The total objective is

$$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$6

The four terms are: $$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$7

$$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$8

$$\mathbf{v}^{n+1}=\mathbf{F}_2(\mathbf{u}_1^n)+\mathbf{G}_2(\mathbf{u}_1^n)\mathbf{v}^n+\boldsymbol{\sigma}_2\boldsymbol{\epsilon}_2^n$$9

and

$$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$0

Here $$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$1 is a warm-up period for the filter (Chen et al., 11 Jul 2025).

Training is carried out in two stages. First, the model is fit without the DA term. Then the diagonal noise levels are estimated from one-step residuals: $$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$2

$$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$3

With $$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$4 fixed, the model is retrained using the full objective including $$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$5 (Chen et al., 11 Jul 2025).

The original CGKN implementation uses fully connected networks for the one-dimensional benchmark systems and a convolutional autoencoder for the $$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$6-D Navier–Stokes case. The continuous-time predecessor emphasizes that DA performance is not an afterthought placed on top of a black-box predictor; it is incorporated directly into the optimization of the model, so the learned latent space is shaped simultaneously by forecast fidelity and inferential tractability (Chen et al., 2024).

5. Relation to Koopman modeling, stochastic lifting, and control

CGKN occupies a specific position within the broader Koopman literature. Standard deep Koopman models seek a lifting $$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$7 such that $$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$8, often with a decoder back to the observation space; this has been demonstrated for video data, discrete spectra, and mode extraction, with explicit latent linear evolution as the architectural constraint (Leask et al., 2020). CGKN keeps the encoder–linear-latent–decoder pattern but relaxes global linearity to conditional linearity, replacing a fixed $$\mathbf{u}^{n+1}=\mathcal{G}(\mathbf{u}^n), \qquad \mathbf{u}^n=(\mathbf{u}_1^n,\mathbf{u}_2^n),$$9 by state-dependent operators $\mathbf{u}_1^n$0 and $\mathbf{u}_1^n$1 acting on the latent state (Chen et al., 11 Jul 2025).

For nonlinear systems with inputs, exact discrete-time Koopman lifting generally does not preserve a constant LTI input matrix. Instead, the exact lifted form is

$\mathbf{u}_1^n$2

or, in scheduling form,

$\mathbf{u}_1^n$3

so the natural object is an LPV-like Koopman model rather than a purely LTI one (Iacob et al., 2022). For control-affine systems, a closely related lifted backbone is the bilinear model

$\mathbf{u}_1^n$4

for which controllability analysis and CLF/LMI-based stabilization have been developed in discrete time (Sinha et al., 2022). These results clarify that conditional or input-dependent Koopman operators are structurally natural rather than ad hoc.

Several neighboring probabilistic formulations illuminate CGKN from different directions. A discrete-time CGKN can be understood as a neural, conditional-Gaussian analogue of Koopman-Equivariant Gaussian Processes, which also place linear time evolution in Koopman features at the center and obtain closed-form multistep Gaussian trajectory distributions through linear propagation of uncertainty (Bevanda et al., 10 Feb 2025). The Stochastic Adversarial Koopman model likewise uses Gaussian latent states, discrete-time Koopman evolution for mean and log-standard deviation, and parameter-conditioned Koopman matrices, but its emphasis is adversarial training and latent sequence prediction rather than analytic DA (Balakrishnan et al., 2021). Physics-informed Koopman networks add generator-based or discrete residual constraints to learn Koopman eigenfunctions with fewer trajectories, suggesting a route to hybrid CGKN training when governing equations or operators are partially known (Liu et al., 2022).

A plausible implication is that control-theoretic Koopman tools such as observability and controllability Gramians, balancing, and balanced truncation can serve as principled latent-dimension diagnostics for CGKN-type models, since they quantify which lifted coordinates matter most for input-output behavior in discrete time (Yeung et al., 2017). At the operator-theoretic end, separate state and input liftings into RKHSs yield an exact bilinear representation

$\mathbf{u}_1^n$5

which suggests a rigorous foundation for conditional Gaussian control extensions of CGKN (Morris et al., 9 Jun 2026).

6. Structured variants, empirical performance, and limitations

The most substantial structured extension is the Lagrangian conditional Gaussian Koopman network (LaCGKN), developed for inferring Eulerian flow fields from moving tracer observations. In that setting, the observed variables are tracer positions and the hidden variables are Eulerian flow fields. LaCGKN preserves the conditional Gaussian latent form but modifies the observation model through tracer homogenization, Fourier positional encoding, and an SVD-inspired low-rank parameterization

$\mathbf{u}_1^n$6

with $\mathbf{u}_1^n$7, reducing parameter count from $\mathbf{u}_1^n$8 to $\mathbf{u}_1^n$9 while maintaining an expressive latent transition operator (Wang et al., 14 Mar 2026).

In the two-layer quasi-geostrophic example of LaCGKN, the reported one-step flow RMSE is $\mathbf{u}_2^n$0 for the high-capacity LaCGKN$\mathbf{u}_2^n$1, compared with $\mathbf{u}_2^n$2 for a DNN(tracer)+CNN(flow) baseline and $\mathbf{u}_2^n$3 for persistence. For posterior flow-field RMSE under DA, LaCGKN reports $\mathbf{u}_2^n$4, compared with $\mathbf{u}_2^n$5 for EnKF, $\mathbf{u}_2^n$6 for optimal interpolation, and $\mathbf{u}_2^n$7 for climatology; the method is also reported as about $\mathbf{u}_2^n$8 faster than parallelized EnKF on the same CPU (Wang et al., 14 Mar 2026).

For the directly discrete-time PDE benchmarks, the reported CGKN errors are as follows (Chen et al., 11 Jul 2025):

System	CGKN forecast MSE	CGKN DA MSE
Viscous Burgers’ equation	$\mathbf{u}_2^n$9	$$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$0
Kuramoto–Sivashinsky equation	$$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$1	$$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$2
$$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$3-D Navier–Stokes equations	$$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$4	$$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$5

These values are accompanied by strong comparisons against both pure forecast surrogates and model-based DA. In the same experiments, FNO attains forecast MSE $$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$6, $$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$7, and $$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$8 on the three systems, while EnKF attains DA MSE $$\mathbf{u}_1^{n+1}=\mathcal{G}_1(\mathbf{u}_1^n,\mathbf{u}_2^n), \qquad \mathbf{u}_2^{n+1}=\mathcal{G}_2(\mathbf{u}_1^n,\mathbf{u}_2^n).$$9, $\mathbf{v}^n$0, and $\mathbf{v}^n$1, respectively. The principal empirical conclusion is therefore not that CGKN dominates all forecast models, but that it attains forecast skill close to state-of-the-art SciML surrogates while simultaneously providing efficient, accurate DA within the same learned architecture (Chen et al., 11 Jul 2025).

The main limitations are structural. The model assumes additive Gaussian white noises in the latent-form surrogate, and it assumes that the hidden dynamics can be represented as conditionally linear in a suitable latent space. Training in current formulations requires offline full-state data for supervision of the hidden state or its encoding. Because $\mathbf{v}^n$2 is nonlinear, physical-space posterior covariance is not analytic and is commonly approximated through residual-based uncertainty networks rather than exact propagation. Finally, latent covariance updates scale cubically in latent dimension in the generic dense case, which motivates low-rank, sparse, or local parameterizations in high-dimensional settings (Chen et al., 2024, Chen et al., 11 Jul 2025).

In this sense, the discrete-time CGKN is best understood not as a generic replacement for deep sequence models, but as a structure-preserving compromise: it forgoes full black-box flexibility in exchange for conditional Gaussian tractability, analytic DA, and a Koopman-style latent organization that remains effective on shocks, spatiotemporal chaos, turbulence, and Lagrangian observation problems (Chen et al., 11 Jul 2025, Wang et al., 14 Mar 2026).